{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "W1_Tutorial2.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernel": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.8"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/CIS-522/course-content/blob/main/tutorials/W01_AlphaZero/W1_Tutorial2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ML2DwqwkVwfo"
      },
      "source": [
        "# CIS-522 Week 1 Part 2\n",
        "# AlphaZero\n",
        "\n",
        "__Instructor:__ Konrad Kording\n",
        "\n",
        "__Content creators:__ Ari Benjamin, Richard Lange, Byron Galbraith\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A9FzYHcnYX4z"
      },
      "source": [
        "---\n",
        "# Today's agenda\n",
        "\n",
        "Today, we will extend what we learned last time all the way to a full AlphaZero (or at least, a tiny one that will run in a Colab).\n",
        "\n",
        "Here's the agenda:\n",
        "1. Review last tutorial\n",
        "2. Introduce the idea of policy and training a policy net\n",
        "3. Talk about planning, tree search, and Monte Carlo Tree Search (MCTS)\n",
        "4. Get AlphaZero!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_MpOtSC-bsnH",
        "cellView": "form"
      },
      "source": [
        "#@markdown What is your Pennkey and pod? (text, not numbers, e.g. bfranklin)\n",
        "my_pennkey = '' #@param {type:\"string\"}\n",
        "my_pod = 'Select' #@param ['Select', 'euclidean-wombat', 'sublime-newt', 'buoyant-unicorn', 'lackadaisical-manatee','indelible-stingray','superfluous-lyrebird','discreet-reindeer','quizzical-goldfish','ubiquitous-cheetah','nonchalant-crocodile','fashionable-lemur','spiffy-eagle','electric-emu','quotidian-lion','astute-jellyfish', 'quantum-herring']\n",
        "\n",
        "# start timing\n",
        "import time\n",
        "try:t0;\n",
        "except NameError: t0 = time.time()\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cQ9Q74E_YYCd",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Towards AlphaZero\n",
        "param_list = ['my_pennkey', 'my_pod']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "from IPython.display import YouTubeVideo\n",
        "video = YouTubeVideo(id=\"GmAPxuPXGhM\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cYdeumcCTu1Y"
      },
      "source": [
        "---\n",
        "# Setup\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hpZ6TBJq554D"
      },
      "source": [
        "### Ensure you're running a GPU notebook.\n",
        "\n",
        "\n",
        "From \"Runtime\" in the drop-down menu above, click \"Change runtime type\". Ensure that \"Hardware Accelerator\" says \"GPU\"."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BqLBTC-ZWpTD"
      },
      "source": [
        "# Imports\n",
        "import math\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "from torch import nn \n",
        "from numba import cuda, float32, int32\n",
        "import numpy as np\n",
        "import os\n",
        "from matplotlib import pyplot as plt\n",
        "import matplotlib.patches as patches\n",
        "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
        "import IPython\n",
        "from IPython.display import IFrame\n",
        "from google.colab import output\n",
        "from tqdm.notebook import tqdm\n",
        "from typing import List\n",
        "import time\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Un34UwDpUkre",
        "cellView": "form"
      },
      "source": [
        "# @title Figure Settings\n",
        "%config InlineBackend.figure_format = 'retina'\n",
        "%matplotlib inline \n",
        "\n",
        "fig_w, fig_h = (8, 6)\n",
        "plt.rcParams.update({'figure.figsize': (fig_w, fig_h)})\n",
        "\n",
        "plt.style.use(\"https://raw.githubusercontent.com/NeuromatchAcademy/course-content/master/nma.mplstyle\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jCBRweGLWp5w"
      },
      "source": [
        "## Review and Recreate\n",
        "To start, we've copied over much of the code from last time (`OthelloGame`, `EpsilonGreedyAgent`, `ai_vs_ai`, etc). Take a moment to make sure everything here looks familiar."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bxty7aS9N-V0",
        "cellView": "form"
      },
      "source": [
        "#@markdown Recreate OthelloGame and InteractiveOthelloGame\n",
        "\n",
        "# Environment config so that numba knows where to find CUDA libraries\n",
        "os.environ['NUMBAPRO_LIBDEVICE'] = \"/usr/local/cuda-10.0/nvvm/libdevice\"\n",
        "os.environ['NUMBAPRO_NVVM'] = \"/usr/local/cuda-10.0/nvvm/lib64/libnvvm.so\"\n",
        "\n",
        "# 1st dimension is row, 2nd is col. Top-left corner of board is (0,0). 8 vector\n",
        "# directions out from each point are:\n",
        "RAYS = np.array([[0, 1],   # east\n",
        "                 [0, -1],  # west\n",
        "                 [1, 0],   # south\n",
        "                 [-1, 0],  # north\n",
        "                 [1, 1],   # southeast\n",
        "                 [1, -1],  # southwest\n",
        "                 [-1, 1],  # northeast\n",
        "                 [-1, -1]],# northwest\n",
        "                dtype=np.int32)\n",
        "\n",
        "# ======================================\n",
        "# === BEGIN CUDA-ACCELERATED HELPERS ===\n",
        "# ======================================\n",
        "\n",
        "def torch2cuda(arr:torch.Tensor) -> cuda.cudadrv.devicearray.DeviceNDArray:\n",
        "    return cuda.as_cuda_array(arr)\n",
        "\n",
        "def cuda2torch(arr:cuda.cudadrv.devicearray.DeviceNDArray) -> torch.Tensor:\n",
        "    return torch.as_tensor(arr, device='cuda')\n",
        "\n",
        "@cuda.jit(device=True)\n",
        "def cuda_cast_rays(board, y0, x0, ray, player):\n",
        "    \"\"\"Helper CUDA kernel that searches out from row, column coordinate (y0,x0)\n",
        "    in the direction specified by \"ray\". Returns True iff there is an unbroken\n",
        "    line of opponent stones along the ray starting at (y0,x0), ending with a\n",
        "    stone belonging to \"player\".\n",
        "    \n",
        "    The device=True decorator means that this may only be called from other CUDA\n",
        "    kernels (see below).\n",
        "    \"\"\"\n",
        "    n = board.shape[1]\n",
        "    opponent = 3 - player\n",
        "    y, x = y0+ray[0], x0+ray[1]\n",
        "\n",
        "    # The ray must begin with an opponent stone 1 space away\n",
        "    if x < 0 or y < 0 or x >= n or y >= n or board[y, x] != opponent:\n",
        "        return False\n",
        "    \n",
        "    # Keep looking until edge of the board or we find an open space or we find\n",
        "    # a player-owned stone\n",
        "    y, x = y+ray[0], x+ray[1]\n",
        "    while x >= 0 and y >= 0 and x < n and y < n:\n",
        "        # Ray terminates on an open space. Not a valid move.\n",
        "        if board[y, x] == 0:\n",
        "            return False\n",
        "        \n",
        "        # Ray terminates on player. This is valid!\n",
        "        if board[y, x] == player:    \n",
        "            return True\n",
        "        \n",
        "        y, x = y+ray[0], x+ray[1]\n",
        "        \n",
        "    # Reached end of the board. Not a valid move.\n",
        "    return False\n",
        "\n",
        "@cuda.jit\n",
        "def cuda_get_valid_actions(boards, player, valid_actions):\n",
        "    \"\"\"Helper CUDA kernel that calls cuda_cast_rays to find all legal moves\n",
        "    in parallel across many boards and positions. Results are written into the\n",
        "    array \"valid_actions\" which must be the same shape as \"boards\".\n",
        "\n",
        "    Decorating with @cuda.jit means we can call it from python by passing cuda\n",
        "    device array objects for \"board\" and \"valid_actions\".\n",
        "\n",
        "    CUDA is structured by 'blocks' and 'threads'. Each block contains multiple\n",
        "    threads. It is the caller's job to say how many of each and their effective\n",
        "    size. For instance, cuda_get_valid_actions[4, (8,8)](...) is the syntax for\n",
        "    \"4 blocks, each of which has 8 'x' and 8 'y' threads\".\n",
        "\n",
        "    We use a separate block per board, and each thread evaluates a single board\n",
        "    coordinate.\n",
        "    \"\"\"\n",
        "    # From numba documentation: const.array_like copies the given array into\n",
        "    # constant GPU memory *at compile time*. This means the copy only happens\n",
        "    # the first time this function is called.\n",
        "    rays = cuda.const.array_like(RAYS)\n",
        "\n",
        "    # We access boards through the 1d block index (boardIdx.x) and positions\n",
        "    # within the board through the thread index x and y.\n",
        "    tx = cuda.threadIdx.x\n",
        "    ty = cuda.threadIdx.y\n",
        "    bx = cuda.blockIdx.x\n",
        "    \n",
        "    # If a space is occupied, it's an illegal move and we're done. Store a zero\n",
        "    # at coordinate (ty, tx) on board bx.\n",
        "    if boards[bx, ty, tx] != 0:\n",
        "        valid_actions[bx, ty, tx] = 0\n",
        "        return\n",
        "\n",
        "    # Try out all 8 ray directions. If any one is a hit, set return value to 1\n",
        "    # and break.\n",
        "    for i in range(8):\n",
        "        hit = cuda_cast_rays(boards[bx], ty, tx, rays[i], player)\n",
        "        if hit:\n",
        "            valid_actions[bx, ty, tx] = 1\n",
        "            return\n",
        "    \n",
        "    # We tried all of the rays and none were legal capturing moves.\n",
        "    valid_actions[bx, ty, tx] = 0\n",
        "\n",
        "@cuda.jit\n",
        "def cuda_step(boards, actions, player):\n",
        "    \"\"\"Helper CUDA kernel for playing actions across multiple boards in parallel.\n",
        "\n",
        "    Boards must be [num_games, n, n], actions must be [num_games, 2] and player\n",
        "    must be 1 or 2. Actions are (row, col) indices of the player's stone.\n",
        "\n",
        "    WARNING: no legality checks are performed here! It is the caller's\n",
        "    responsibility to ensure that all moves are legal.\n",
        "\n",
        "    Boards are indexed by block x and the 8 ray directions by thread x\n",
        "\n",
        "    If either of actions[b,:] is -1, treated as pass.\n",
        "    \"\"\"\n",
        "    # As in cuda_get_valid_actions, this copy happens once at compile time\n",
        "    rays = cuda.const.array_like(RAYS)\n",
        "\n",
        "    # Grab thread and block index from cuda context\n",
        "    tx = cuda.threadIdx.x\n",
        "    bx = cuda.blockIdx.x  \n",
        "\n",
        "    # Unpack the row, col coordinate of the action\n",
        "    act_y, act_x = actions[bx]\n",
        "\n",
        "    # If -1, the player passed and there is nothing to do.\n",
        "    if act_x == -1 or act_y == -1:\n",
        "        return\n",
        "\n",
        "    # Search a different ray direction on each of 8 threads. ASSUMES the move\n",
        "    # is legal.\n",
        "    is_hit = cuda_cast_rays(boards[bx], act_y, act_x, rays[tx], player)\n",
        "    if is_hit:\n",
        "        dy, dx = rays[tx]\n",
        "        opponent = 3 - player\n",
        "        y, x = act_y+dy, act_x+dx\n",
        "        while boards[bx, y, x] == opponent:\n",
        "            # Modify the board, flipping opponent to player\n",
        "            boards[bx, y, x] = player\n",
        "            y, x = y+dy, x+dx\n",
        "    \n",
        "    # Only need to do this once, so just let thread #0 handle it: update\n",
        "    # the board at location of the action\n",
        "    if tx == 0:\n",
        "        boards[bx, act_y, act_x] = player\n",
        "\n",
        "@torch.jit.script\n",
        "def zobrist_hash(boards:torch.Tensor, player:torch.Tensor, table:torch.Tensor, out:torch.Tensor):\n",
        "    n_games, n, _ = boards.size()\n",
        "    flat_boards = boards.reshape(n_games, n*n)\n",
        "    bit_strings = table[torch.arange(n*n), flat_boards.long()]\n",
        "    # Reset output values to 0\n",
        "    out.fill_(0)\n",
        "    # We have to loop manually because bitwise_xor has no dimension\n",
        "    # argument (see https://github.com/pytorch/pytorch/issues/35641)\n",
        "    for i in range(n*n):\n",
        "        torch.bitwise_xor(out, bit_strings[:,i], out=out)\n",
        "    # Just add the current player - ok that this isn't random bitstrings because\n",
        "    # just a single integer added per game.\n",
        "    torch.add(out, player, out=out)\n",
        "    return out\n",
        "\n",
        "# ====================================\n",
        "# === BEGIN OTHELLOGAME DEFINITION ===\n",
        "# ====================================\n",
        "\n",
        "class OthelloGame(object):\n",
        "    PLAYER1 = 1\n",
        "    PLAYER2 = 2\n",
        "    ACTIVE = 0\n",
        "    OVER = 1\n",
        "    PASS = (-1, -1)\n",
        "    # Unicode characters for dark and light circles if printing to console.\n",
        "    # By convention, player 1 is dark and 2 is light.\n",
        "    GAME_SYMBOLS = ['_', '\\u25cf', '\\u25cb']\n",
        "\n",
        "    def __init__(self, n_games=1, n=8):\n",
        "        \"\"\"Create a new batch of n_games othello games each with board size [n,n]\n",
        "\n",
        "        The game states are managed on the GPU with the help of CUDA kernels.\n",
        "\n",
        "        We expect all n games to update in lockstep, so that there is only one\n",
        "        'current_player' shared by all of them.\n",
        "        \"\"\"\n",
        "        self.n_games = n_games\n",
        "        self.n = n\n",
        "        self.boards = torch.zeros((n_games, n, n), dtype=torch.float32, device='cuda')\n",
        "        # Initial positions: 4 stones in the center alternating color\n",
        "        self.boards[:, [self.n//2, self.n//2-1], [self.n//2, self.n//2-1]] = OthelloGame.PLAYER2\n",
        "        self.boards[:, [self.n//2, self.n//2-1], [self.n//2-1, self.n//2]] = OthelloGame.PLAYER1\n",
        "        # Always start with player 1 (dark stones by convention)\n",
        "        self.current_player = OthelloGame.PLAYER1\n",
        "        # All games are initially active\n",
        "        self.game_status = OthelloGame.ACTIVE*torch.ones(n_games, dtype=torch.int32, device='cuda')\n",
        "        # If both players pass the game is over. Keep track of whether the last\n",
        "        # player passed.\n",
        "        self.last_player_pass = torch.zeros(n_games, dtype=torch.int32, device='cuda')\n",
        "\n",
        "        # _valid_moves is a container for storing which moves are valid. It is\n",
        "        # zero everywhere except where it is legal to play, which is 1.\n",
        "        self._valid_moves = torch.zeros_like(self.boards)\n",
        "\n",
        "        # Create a CUDA device array copy of each GPU tensor. These share GPU\n",
        "        # memory with torch tensors so any changes made in CUDA kernels are seen\n",
        "        # by torch.\n",
        "        self._cuda_boards = torch2cuda(self.boards)\n",
        "        self._cuda_valid = torch2cuda(self._valid_moves)\n",
        "\n",
        "        # Populate valid moves using CUDA-ized algorithm\n",
        "        self.refresh_legal()\n",
        "\n",
        "        # Initialize 'zobrist hash' table for get_uid() function\n",
        "        self._new_zobrist_table()\n",
        "\n",
        "    def step(self, actions:torch.Tensor):\n",
        "        \"\"\"Place a stone for the current player at location action=(row,col)\n",
        "        separately per board. Expected size of 'actions' is [n_games, 2].\n",
        "\n",
        "        Action coordinates of -1 treated as passing.\n",
        "\n",
        "        Actions are ignored for all games that are not in the ACTIVE state.\n",
        "\n",
        "        Note: no legality checking here! It is the caller's responsibility to\n",
        "        ensure 'action' is legal, e.g. by checking that action is in\n",
        "        game.get_available_actions().\n",
        "        \"\"\"\n",
        "        # Ensure actions is a torch tensor and on the GPU. Ensure its shape is\n",
        "        # [n_games, 2] even if the input was, say, just a (row,col) tuple.\n",
        "        actions = torch.as_tensor(actions, device='cuda').view(-1, 2)\n",
        "\n",
        "        # Ensure actions are \"pass\" for games that have completed. This tells\n",
        "        # the cuda_step kernel to ignore these games.\n",
        "        actions[self.game_status == OthelloGame.OVER, :] = -1\n",
        "\n",
        "        # Update board state with call to CUDA kernel cuda_step.\n",
        "        cuda_step[self.n_games, RAYS.shape[0]](self._cuda_boards, torch2cuda(actions), self.current_player)\n",
        "\n",
        "        # Advance to the next player (2->1 and 1->2).\n",
        "        self.current_player = 3-self.current_player\n",
        "\n",
        "        # Count passes - any games with 2 passes in a row is flagged as being over.\n",
        "        is_pass = torch.any(actions == -1, dim=1)\n",
        "        self.last_player_pass[~is_pass] = 0\n",
        "        self.last_player_pass[is_pass] += 1\n",
        "\n",
        "        # Flag games that have finished\n",
        "        is_game_over = self.last_player_pass >= 2\n",
        "        self.game_status[is_game_over] = OthelloGame.OVER\n",
        "\n",
        "        # Refresh legal moves for the next turn\n",
        "        self.refresh_legal()\n",
        "    \n",
        "    def refresh_legal(self):\n",
        "        # Look for valid moves in the updated boards for the new player\n",
        "        cuda_get_valid_actions[self.n_games, (self.n, self.n)](self._cuda_boards, self.current_player, self._cuda_valid)\n",
        "\n",
        "    def copy_state(self, idx=None):\n",
        "        # Grab the minimal amount of state info to be able to return the board\n",
        "        # to its current configuration later with a call to paste_state.\n",
        "        # NOTE: this copy is single-use only; once it is pasted, it will be\n",
        "        # modified in-place!\n",
        "        return {\"boards\": self.boards.clone(),\n",
        "                \"game_status\": self.game_status.clone(),\n",
        "                \"last_player_pass\": self.last_player_pass.clone(),\n",
        "                \"current_player\": self.current_player}\n",
        "\n",
        "    def paste_state(self, state):\n",
        "        self.__dict__.update(state)\n",
        "        self._cuda_boards = torch2cuda(self.boards)\n",
        "        self._cuda_valid = torch2cuda(self._valid_moves)\n",
        "        cuda_get_valid_actions[self.n_games, (self.n, self.n)](self._cuda_boards, self.current_player, self._cuda_valid)\n",
        "\n",
        "    def render(self, mode='human'):\n",
        "        boards = self.boards.cpu()\n",
        "        acts = self.get_available_actions()\n",
        "        for g in range(self.n_games):\n",
        "            if mode == 'text':\n",
        "                print(\"=\"*(self.n-1) + (f\"{OthelloGame.GAME_SYMBOLS[self.current_player]}\")*2 + \"=\"*(self.n-1))\n",
        "                for row in boards[g]:\n",
        "                    print(\" \".join(OthelloGame.GAME_SYMBOLS[v] for v in row))\n",
        "                print(\"=\"*(self.n*2))\n",
        "            else:\n",
        "                # Plot with matplotlib patches. Player 1 is dark and 2 is light.\n",
        "                fig, ax = plt.subplots(figsize=(6,6))\n",
        "                ax.set_aspect('equal')\n",
        "                rect = patches.Rectangle((0,0),self.n,self.n,linewidth=1,edgecolor='k',facecolor='g')\n",
        "                ax.add_patch(rect)\n",
        "                ax.set_xlim(0,self.n)\n",
        "                ax.set_ylim(0,self.n)\n",
        "                ax.set_yticklabels([])\n",
        "                ax.set_xticklabels([])\n",
        "                ax.tick_params(length=0)\n",
        "                ax.grid(which='major', zorder=0, c='k')\n",
        "                for i in range(self.n):\n",
        "                    for j in range(self.n):\n",
        "                        if boards[g,i,j]>0:\n",
        "                            # Draw stones\n",
        "                            c = (0,0,0,1) if boards[g,i,j]==OthelloGame.PLAYER1 else (1,1,1,1)\n",
        "                            circ = patches.Circle((j+.525,self.n-1-i+.5), .4, \n",
        "                                                facecolor=c, linewidth=1,edgecolor='k')\n",
        "                            ax.add_patch(circ)\n",
        "                        elif (i,j) in acts[g]:\n",
        "                            # Draw transluscent stones on legal move positions\n",
        "                            c = (0,0,0,0.25) if self.current_player==OthelloGame.PLAYER1 else (1,1,1,0.25)\n",
        "                            circ = patches.Circle((j+.525,self.n-1-i+.5), .4, facecolor=c, linewidth=0)\n",
        "                            ax.add_patch(circ)\n",
        "                plt.show()\n",
        "\n",
        "    def get_available_actions(self):\n",
        "        \"\"\"Get a list of all legal actions for the current player (passing not included).\n",
        "        \n",
        "        Note: assumes self._valid_moves is up to date!\n",
        "        \"\"\"\n",
        "        acts = [None]*self.n_games\n",
        "        for g in range(self.n_games):\n",
        "            i, j = torch.where(self._valid_moves[g] == 1)\n",
        "            acts[g] = list(zip(i.cpu().numpy(), j.cpu().numpy()))\n",
        "        return acts\n",
        "\n",
        "    def are_games_over(self):\n",
        "        return torch.all(self.game_status == OthelloGame.OVER).item()\n",
        "\n",
        "    def score_games(self):\n",
        "        \"\"\"Compute [n_games, 2] tuple containing \"scores\" for player 1 and 2 in\n",
        "        each column. Assumes game is over.\n",
        "\n",
        "        scores[:,0] is +1 if player 1 won, 0 for draw, or -1 for loss\n",
        "        scores[:,1] is the same from player 2's perspective\n",
        "        \"\"\"\n",
        "        flat_boards = self.boards.view(self.n_games, -1)\n",
        "        n_dark = torch.sum(flat_boards == OthelloGame.PLAYER1, dim=1)\n",
        "        n_lite = torch.sum(flat_boards == OthelloGame.PLAYER2, dim=1)\n",
        "        scores = torch.zeros_like(n_dark)\n",
        "        scores[n_dark > n_lite] = +1\n",
        "        scores[n_dark < n_lite] = -1\n",
        "        return torch.stack([scores,-scores], dim=1)\n",
        "    \n",
        "    def get_uid(self, player=None, out=None):\n",
        "        \"\"\"Return highly-probably-unique identifier for each game. Output is a\n",
        "        tensor of [n_games] int64 values.\n",
        "\n",
        "        Algorithm is Zobrist hashing.\n",
        "        \"\"\"\n",
        "        if player is None:\n",
        "            player = self.current_player * torch.ones(self.n_games, dtype=torch.int64, device='cuda')\n",
        "        elif not isinstance(player, torch.Tensor) or len(player) == 1:\n",
        "            player = player * torch.ones(self.n_games, dtype=torch.int64, device='cuda')\n",
        "        if out is None:\n",
        "            out = torch.zeros(self.n_games, dtype=torch.int64, device='cuda')\n",
        "        return zobrist_hash(self.boards, player, self._hash_board, out)\n",
        "\n",
        "    def _new_zobrist_table(self):\n",
        "        self._hash_board = torch.randint(2**63-1, size=(self.n*self.n, 3), dtype=torch.int64, device='cuda')\n",
        "        # Sanity check that no 2 random strings were identical\n",
        "        assert len(self._hash_board.flatten().unique()) == self.n*self.n*3\n",
        "\n",
        "################################\n",
        "#### INTERACTIVE GAME BOARD ####\n",
        "################################\n",
        "\n",
        "def temporary_info(message, clear=False):\n",
        "    \"\"\"Output overwritable message to the console. If clear=True, overwrite all old messages\n",
        "    \"\"\"\n",
        "    if clear:\n",
        "        output.clear(output_tags=\"temporary-info\")\n",
        "    with output.use_tags(\"temporary-info\"):\n",
        "        print(message)\n",
        "\n",
        "class InteractiveOthelloGame(object):\n",
        "    def __init__(self, game=None, player1=\"human\", player2=\"human\", n=8):\n",
        "        if game is None:\n",
        "            self.game = OthelloGame(n_games=1, n=n)\n",
        "        else:\n",
        "            self.game = game\n",
        "            assert game.n_games == 1, \"Can't handle >1 games interactively!\"\n",
        "        self.players = [player1, player2]\n",
        "\n",
        "    def next_turn(self):\n",
        "        self.available_acts = self.game.get_available_actions()[0]\n",
        "        if self.game.game_status == OthelloGame.ACTIVE:\n",
        "            the_player = self.players[self.game.current_player-1]\n",
        "            temporary_info(f\"Begin Player {self.game.current_player}'s turn\", clear=True)\n",
        "            if the_player == \"human\":\n",
        "                if len(self.available_acts) > 0:\n",
        "                    temporary_info(\"Input move by clicking the board\")\n",
        "                else:\n",
        "                    temporary_info(\"No legal moves! Click anywhere on the board to pass.\")\n",
        "                self.redraw()\n",
        "                self._reregister_click_callback()\n",
        "            else:\n",
        "                err = False\n",
        "                temporary_info(\"Waiting for AI to complete...\")\n",
        "                try:\n",
        "                    ai_action = the_player.select_move(self.game)\n",
        "                    try:\n",
        "                        self.game.step(ai_action)\n",
        "                    except Exception as e:\n",
        "                        err = True\n",
        "                        print(\"Error in game.step() on AI-selected move!\")\n",
        "                        print(e)\n",
        "                except Exception as e:\n",
        "                    err = True\n",
        "                    print(\"Error in AI select_move!\")\n",
        "                    print(e)\n",
        "                if not err:\n",
        "                    self.next_turn()\n",
        "        else:\n",
        "            self.redraw()\n",
        "            temporary_info(\"GAME OVER!\", clear=True)\n",
        "            values = self.game.score_games()\n",
        "            if values[0] == 0:\n",
        "                print(f\"Game ended in a draw\")\n",
        "            else:\n",
        "                winner_id = 1+np.argmax(values)\n",
        "                try:\n",
        "                    name = self.players[winner_id-1].name\n",
        "                except AttributeError:\n",
        "                    name = f\"Player {winner_id}\"\n",
        "                print(f\"{name} is the winner!\")\n",
        "\n",
        "    def handle_user_click(self, x_pix, y_pix, plot_width, plot_height):\n",
        "        if len(self.available_acts) == 0:\n",
        "            # If there is not available action all you can do is pass\n",
        "            self.game.step(OthelloGame.PASS)\n",
        "            self.next_turn()\n",
        "        else:\n",
        "            cell_width = plot_width / self.game.boards.shape[2]\n",
        "            cell_height = plot_height / self.game.boards.shape[1]\n",
        "            cell_x, cell_y = int(x_pix / cell_width), int(y_pix / cell_height)\n",
        "            action = (cell_y, cell_x)\n",
        "            if action in self.available_acts:\n",
        "                self.game.step(action)\n",
        "                self.next_turn()\n",
        "            else:\n",
        "                temporary_info(f\"Available actions are {self.available_acts} but you clicked {action}\")\n",
        "\n",
        "    def _reregister_click_callback(self):\n",
        "        # Inject javascript which will detect a click and invoke a function called 'pass_to_python_handler'\n",
        "        display(IPython.display.Javascript(\"\"\"\n",
        "        var plot_element = document.querySelector(\".output_image\").firstElementChild;\n",
        "        plot_element.onclick = function(event){\n",
        "            google.colab.kernel.invokeFunction(\"pass_to_python_handler\", [event.offsetX, event.offsetY, plot_element.width, plot_element.height], {});\n",
        "        };\n",
        "        \"\"\"))\n",
        "\n",
        "        # Tell colab that when 'pass_to_python_handler' is called in JS, it should\n",
        "        # call self.handle_user_click in python\n",
        "        output.register_callback(\"pass_to_python_handler\", self.handle_user_click)\n",
        "\n",
        "    def redraw(self):\n",
        "        # Clear previous output\n",
        "        output.clear(output_tags='othello-interactive')\n",
        "        # Draw a fresh plot and store the figsize in pixels\n",
        "        with output.use_tags('othello-interactive'):\n",
        "            self.game.render()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-25OdE2WSmyT",
        "cellView": "form"
      },
      "source": [
        "#@markdown Recreate ai_vs_ai and win_statistics\n",
        "\n",
        "def ai_vs_ai(agent1, agent2=None, n_games=1, n=8, randomize_players=False,\n",
        "             print_result=False, progbar=True, callback=None, del_progbar=False):\n",
        "    \"\"\"Pit 2 AI players against each other (or self-play if just 1 given). Return\n",
        "    end-state values and list of board positions.\n",
        "    \"\"\"\n",
        "    agent1.new_game()\n",
        "    if agent2 is None:\n",
        "        agent2 = agent1\n",
        "    else:\n",
        "        agent2.new_game()\n",
        "    game = OthelloGame(n_games=n_games, n=n)\n",
        "    saved_boards = []\n",
        "    callback_data = []\n",
        "    agents = [agent1, agent2]\n",
        "    if randomize_players:\n",
        "        agents = [agent2, agent1]\n",
        "    if progbar:\n",
        "        # Games rarely go beyond n^2-4 moves since passes are rare\n",
        "        progress_bar = tqdm(total=game.n**2-4, desc='# moves per game',\n",
        "                            leave=not del_progbar)\n",
        "    \n",
        "    while torch.any(game.game_status == OthelloGame.ACTIVE):\n",
        "        # Save board state and flag all boards that are done by flooding them with nan\n",
        "        saved_boards.append(game.boards.clone())\n",
        "        saved_boards[-1][game.game_status == OthelloGame.OVER, ...] = float('nan')\n",
        "        \n",
        "        # Get move by current player\n",
        "        act = agents[game.current_player-1].select_move(game)\n",
        "        \n",
        "        # Do callback after AI decision but before updating sate\n",
        "        if callback is not None:\n",
        "            callback_data.append(callback(game, agents))\n",
        "        \n",
        "        # Update state\n",
        "        game.step(act)\n",
        "        \n",
        "        # Update progress bar\n",
        "        if progbar: progress_bar.update(1)\n",
        "    \n",
        "    if progbar:\n",
        "        progress_bar.total = progress_bar.n\n",
        "        progress_bar.close()\n",
        "    values = game.score_games()\n",
        "    if print_result:\n",
        "        for g in n_games:\n",
        "            print(f\"Game {g+1} of {n_games}:\", end=\"\\t\")\n",
        "            if values[0] == 0:\n",
        "                print(\"Draw\")\n",
        "            else:\n",
        "                print(f\"{agents[torch.argmax(values)].name} (Player {torch.argmax(values)+1}) is the winner!\")\n",
        "    # Return end-game value according to player 1 and 2, as well as a list of\n",
        "    # all board states throughout play\n",
        "    return values, saved_boards, callback_data\n",
        "\n",
        "def outcome2color(outcome, alpha=1):\n",
        "    \"\"\"Return RGB color for given outcome (-1, 0, +1)\n",
        "    \"\"\"\n",
        "    if outcome == +1:\n",
        "        return (0, 1, 0, alpha) # Win as green\n",
        "    elif outcome == -1:\n",
        "        return (1, 0, 0, alpha) # Loss as red\n",
        "    else:\n",
        "        return (0, 0, 1, alpha) # Draw as blue\n",
        "\n",
        "def win_statistics(agent1, agent2, games_per_side=100, bar_plot=True):\n",
        "    \"\"\"Play agent1 vs agent2 many times to get some statistics. At the end, show\n",
        "    statistics with a bar plot and return wins/draws/losses for agent1. Runs\n",
        "    'games_per_side' matches with agent1 as player 1 and another 'games_per_side'\n",
        "    matches with agent1 as player 2.\n",
        "\n",
        "    The returned 'win_info' table is size 2x3. Row 0 counts agent1 as player 1.\n",
        "    Row 1 counts agent1 as player 2. Col 0 is wins for agent1, col 1 is draws,\n",
        "    and col 2 is losses\n",
        "    \"\"\"\n",
        "    win_info = np.zeros((1, 3))\n",
        "    # Play agent1 as player 1\n",
        "    outcome1, _, _ = ai_vs_ai(agent1, agent2, n_games=games_per_side//2, \n",
        "                          n=6)\n",
        "    # Play agent1 as player 2\n",
        "    outcome2, _, _ = ai_vs_ai(agent2, agent1, n_games=games_per_side//2, \n",
        "                          n=6)\n",
        "    outcome = torch.cat((outcome1,outcome2),dim=0) \n",
        "    for g in range(games_per_side):\n",
        "        col_idx = 1-outcome[g,0] # Map agent1 outcome (-1,0,1) to column index (2,1,0)\n",
        "        win_info[0, col_idx] += 1\n",
        "    # (Maybe) plot\n",
        "    if bar_plot:\n",
        "        # Approximate standard error in each probability estimate (it would be\n",
        "        # more correct to do confidence intervals on the binomial distribution\n",
        "        # for each of win/draw/loss, but this is fine for a large # of games)\n",
        "        win_prob = win_info / games_per_side\n",
        "        win_se = np.sqrt(win_prob*(1-win_prob)/games_per_side)\n",
        "        plt.figure(figsize=(4,6))\n",
        "        plt.bar([0], win_prob[:,0], yerr=win_se[:,0], color=outcome2color(1, .8), label=\"wins\")\n",
        "        plt.bar([0], win_prob[:,1], yerr=win_se[:,1], bottom=win_prob[:,0], color=outcome2color(0, 0.8))\n",
        "        plt.bar([0], win_prob[:,2], yerr=win_se[:,2], bottom=win_prob[:,:2].sum(axis=1), color=outcome2color(-1, 0.8))\n",
        "        plt.xticks([])\n",
        "        plt.yticks(np.linspace(0,1,11))\n",
        "        plt.ylabel(\"Fraction games won\")\n",
        "        plt.title(f\"{agent1.name} vs {agent2.name}\")\n",
        "        plt.legend([f\"{agent1.name} wins\", \"draws\", f\"{agent1.name} loses\"])\n",
        "        plt.grid()\n",
        "        plt.show()\n",
        "    return win_info"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Fo8N2lqxSdR_",
        "cellView": "form"
      },
      "source": [
        "#@markdown Recreate epsilon-greedy player and hand-coded value function\n",
        "\n",
        "def my_board_value(board_states:torch.Tensor, whoami:torch.Tensor) -> torch.Tensor:\n",
        "    \"\"\"Like my_board_value from last time but now 'whoami' is a [n_games] size tensor\n",
        "    \"\"\"\n",
        "    # An easy option is the proportion of tiles, rescaled to [-1,1]. \n",
        "    # Another could be # of available moves for your opponent.\n",
        "\n",
        "    n_games, n, _ = board_states.size()\n",
        "    values = torch.zeros(n_games)\n",
        "\n",
        "    # to-do: have students write their own code here.\n",
        "    n1 = (board_states == OthelloGame.PLAYER1).view(n_games, n*n).sum(dim=1)\n",
        "    n2 = (board_states == OthelloGame.PLAYER2).view(n_games, n*n).sum(dim=1)\n",
        "\n",
        "    frac1 = n1/(n1+n2)\n",
        "    values = frac1*2-1\n",
        "    values[whoami == OthelloGame.PLAYER2] = -values[whoami == OthelloGame.PLAYER2]\n",
        "\n",
        "    return values\n",
        "\n",
        "class ArtificialPlayer(object):\n",
        "    \"\"\"A base class for agents that play Othello. Sub-classes must provide a\n",
        "    select_move method which takes in an OthelloGame and returns an action\n",
        "    (in other words, returns a tuple (row,col) to place a stone or None to pass)\n",
        "    \"\"\"\n",
        "    def __init__(self, name=\"AI Agent\"):\n",
        "        self.name = name\n",
        "\n",
        "    def select_move(self, game:OthelloGame):\n",
        "        raise NotImplementedError(\"Must subclass with a playing rule\")\n",
        "    \n",
        "    def new_game(self):\n",
        "        \"\"\"Reset agent state. By default, do nothing...\n",
        "        \"\"\"\n",
        "        pass\n",
        "\n",
        "class EpsilonGreedyAgent(ArtificialPlayer):\n",
        "    def __init__(self, epsilon=0.05, name=None):\n",
        "        super(EpsilonGreedyAgent, self).__init__(name if name is not None else f\"Epsilon={epsilon}\")\n",
        "        self.epsilon = epsilon\n",
        "\n",
        "    def evaluate(self, game:OthelloGame, whoami:torch.Tensor):\n",
        "        \"\"\"\n",
        "        Evaluation function used in \"select_move\" below. Subclasses may override\n",
        "        evaluate() to create epsilon-greedy players with a different value function.\n",
        "        \"\"\"\n",
        "        # Default to the hand-coded heuristic from above\n",
        "        return my_board_value(game.boards, whoami)\n",
        "    \n",
        "    def batch_evaluate(self, game:OthelloGame, whoami:torch.Tensor):\n",
        "        \"\"\"Call evaluate() once per available action. Returns three things:\n",
        "            - 'values' a [n_games, n_actions] tensor containing values estimated\n",
        "              for each action per game.\n",
        "            - 'actions_tensor' a [n_games, n_actions, 2] tensor containing the\n",
        "              (row,col) coordinates of corresponding actions.\n",
        "            - 'actions_list' the output of game.get_available_actions\n",
        "\n",
        "        Handles the case of game.n_games>1 somewhat efficiently. If different\n",
        "        games have different #s of actions, then n_actions is the maximum number\n",
        "        of actions in any of the games. 'actions' is padded with -1 on games\n",
        "        that have fewer actions (since -1 indicates \"pass\") and 'values' is set\n",
        "        to -inf on all pass moves.\n",
        "        \"\"\"\n",
        "        # actions_list is a list of lists of tuples. One list per game.\n",
        "        actions_list = game.get_available_actions()\n",
        "        num_actions_each_game = torch.as_tensor([len(acts) for acts in actions_list])\n",
        "        max_actions = num_actions_each_game.max()\n",
        "        # Allocate values output\n",
        "        values = float('-inf')*torch.ones(game.n_games, max_actions)\n",
        "        # Create a tensor of [n_games, max_actions, 2] that contains all available\n",
        "        # actions for all games, padded with -1s (pass move) where games have\n",
        "        # fewer than max_actions possible actions\n",
        "        actions_tensor = -torch.ones(game.n_games, max_actions, 2, dtype=torch.int32)\n",
        "        for g in range(game.n_games):\n",
        "            if num_actions_each_game[g] == 0:\n",
        "                continue\n",
        "            actions_tensor[g, :num_actions_each_game[g], :] = torch.as_tensor(actions_list[g], dtype=torch.int32)\n",
        "        \n",
        "        # Evaluate each action, simultaneously for all games\n",
        "        for i in range(max_actions):\n",
        "            state = game.copy_state()\n",
        "            # note: games with fewer actions than max_actions will pass here - and switch current players\n",
        "            # but we calculate value only with whoami\n",
        "            game.step(actions_tensor[:,i,:])\n",
        "            values[:,i] = self.evaluate(game, whoami)\n",
        "            values[num_actions_each_game < i+1,i] = float('-inf')\n",
        "            game.paste_state(state)\n",
        "        return values, actions_tensor, actions_list, num_actions_each_game\n",
        "\n",
        "    def select_move(self, game:OthelloGame):\n",
        "        # Evaluate all actions in all games.\n",
        "        all_values, actions_tensor, actions_list,_ = self.batch_evaluate(game, game.current_player*torch.ones_like(game.boards[:,0,0]))\n",
        "        # Initially select the best action in each game.\n",
        "        chosen_action = torch.zeros(game.n_games, 2, dtype=torch.int32)\n",
        "        for g in range(game.n_games):\n",
        "            # If there are no moves available, pass\n",
        "            if len(actions_list[g]) == 0:\n",
        "                chosen_action[g, :] = -1\n",
        "                continue\n",
        "            if np.random.rand() < self.epsilon:\n",
        "                acts = actions_list[g]\n",
        "                random_choice = acts[np.random.randint(len(acts))]\n",
        "                chosen_action[g, :] = torch.as_tensor(random_choice, dtype=torch.int32)\n",
        "            else:\n",
        "                idx = torch.argmax(all_values[g])\n",
        "                chosen_action[g, :] = actions_tensor[g, idx, :]\n",
        "\n",
        "        return chosen_action\n",
        "\n",
        "    def select_move_vectorized(self, game:OthelloGame):\n",
        "        # Evaluate all actions in all games.\n",
        "        all_values, actions_tensor, _, num_actions_each_game = self.batch_evaluate(game, game.current_player*torch.ones_like(game.boards[:,0,0]))\n",
        "\n",
        "        # If there are no moves available for anyone, pass\n",
        "        # (pass is indicated by the (-1,-1) action)\n",
        "        if num_actions_each_game.max()==0:\n",
        "            return -torch.ones(game.n_games, 2, dtype=torch.int32)\n",
        "\n",
        "        if np.random.rand() < self.epsilon:\n",
        "            selection = (num_actions_each_game*torch.rand(game.n_games)).to(int)\n",
        "            chosen_action = actions_tensor[range(game.n_games),selection]\n",
        "        else:\n",
        "            selection = torch.argmax(all_values, dim=1)\n",
        "            chosen_action = actions_tensor[range(game.n_games),selection]\n",
        "\n",
        "        # If there are no moves available for some games, pass those\n",
        "        # (pass is indicated by the (-1,-1) action)\n",
        "        chosen_action[num_actions_each_game==0] = -1\n",
        "\n",
        "        return chosen_action"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uVn7Hx1eZHFG",
        "cellView": "form"
      },
      "source": [
        "#@markdown Recreate NN-training helpers from last time\n",
        "\n",
        "# can you make sense of this list of symmetries?\n",
        "SYMMETRIES = [lambda x: x,\n",
        "              lambda x: torch.rot90(x, 1, dims=(1, 2)),\n",
        "              lambda x: torch.rot90(x, 2, dims=(1, 2)),\n",
        "              lambda x: torch.rot90(x, 3, dims=(1, 2)),\n",
        "              lambda x: torch.flip(x, dims=(1,)),\n",
        "              lambda x: torch.flip(x, dims=(2,)),\n",
        "              lambda x: torch.transpose(x, 1, 2)]\n",
        "\n",
        "def boards2nn(board_states) -> torch.Tensor:\n",
        "    \"\"\"Canonical preprocessing, taking in a set of boards and outputting NN inputs.\n",
        "    \"\"\"\n",
        "    if isinstance(board_states, list):\n",
        "        board_states = torch.cat(board_states, dim=0)\n",
        "    else:\n",
        "        board_states = board_states.clone() # Copy so we don't affect the original\n",
        "    # Convert from [0,1,2] to [0,1,-1] (player 2 represented with -1)\n",
        "    board_states[torch.where(board_states == 2)] = -1\n",
        "    return board_states\n",
        "\n",
        "def prepare_examples(player1_value:torch.Tensor, board_states:List[torch.Tensor]):\n",
        "    \"\"\"Take raw game observations (final value from P1 perspective, list of\n",
        "    boards) and prepare a batch of training examples for the NN. Does 3 things\n",
        "    to the input:\n",
        "    1. Concatenates observations together\n",
        "    2. Adds \"inverse\" of all the inputs to account for player 2's perspective\n",
        "    3. Adds symmetries\n",
        "\n",
        "    In total, this makes the effective # examples 16 times as big as the number \n",
        "    of input positions (8 symmetries from each of 2 players' perspectives)\n",
        "    \n",
        "    player1_value is size [n_games]\n",
        "    board_states is size [max_num_moves]. Each entry contains a [n_games,n,n] board\n",
        "    state for each game at that move.\n",
        "    \"\"\"\n",
        "    n_moves = len(board_states)\n",
        "    n_games, n, _ = board_states[0].size()\n",
        "    # Concatenate and preprocess board states. Result is size [n_games*n_moves,n,n]\n",
        "    # ordered like [game0turn0, game1turn0, ... gameNturn0, game0turn1, game1turn1, ... ]\n",
        "    board_states = boards2nn(board_states)\n",
        "    # Repeat all the player-1 values once per move. The input value are already\n",
        "    # ordered like [game0, game1, ..., gameN] so to match the board order we\n",
        "    # simply repeat() the values\n",
        "    outcomes = player1_value.repeat(n_moves).float()\n",
        "    # It may be the case that not all games ran for n_moves. These are indicated\n",
        "    # by nan values in the board state (see ai_vs_ai)\n",
        "    drop_states = torch.isnan(board_states[:,0,0])\n",
        "    board_states = board_states[~drop_states]\n",
        "    outcomes = outcomes[~drop_states]\n",
        "    # Add all symmetries (increases effective data 8-fold)\n",
        "    outcomes = outcomes.repeat(len(SYMMETRIES))\n",
        "    board_states = torch.cat([sym(board_states) for sym in SYMMETRIES], dim=0)\n",
        "    # Add in examples from player 2's perspective by flipping signs on both\n",
        "    # value and on the board state itself\n",
        "    outcomes = torch.cat([outcomes, -outcomes], dim=0)\n",
        "    board_states = torch.cat([board_states, -board_states], dim=0)\n",
        "    return outcomes, board_states\n",
        "\n",
        "def plot_value_predictions(value_net, agent, n_games=100, n=8, title=None):\n",
        "    value_net.eval() # Ensure 'evaluation' rather than 'training' mode\n",
        "\n",
        "    plt.figure(figsize=(8,4))\n",
        "    ax1 = plt.subplot(1,2,1)\n",
        "    ax2 = plt.subplot(1,2,2)\n",
        "    outcomes, boards, _ = ai_vs_ai(agent, n_games=n_games, progbar=True, n=n)\n",
        "    board_nn = boards2nn(boards)\n",
        "    # Note: some of boards_nn will be nan if the games ended early, in which\n",
        "    # case we don't want to plot them anyway. Force those outputs to nan too.\n",
        "    predicted_values_1 = value_net(board_nn)\n",
        "    predicted_values_2 = value_net(-board_nn)\n",
        "    predicted_values_1[torch.isnan(board_nn[:,0,0])] = float('nan')\n",
        "    predicted_values_2[torch.isnan(board_nn[:,0,0])] = float('nan')\n",
        "    for i in range(n_games):\n",
        "        vals_1_this_game = predicted_values_1[i::n_games]\n",
        "        vals_2_this_game = predicted_values_2[i::n_games]\n",
        "        ax1.plot(vals_1_this_game.detach().cpu(), color=outcome2color(outcomes[i,0].item(), 0.1))\n",
        "        ax2.plot(vals_2_this_game.detach().cpu(), color=outcome2color(outcomes[i,1].item(), 0.1))\n",
        "    for i, ax in enumerate([ax1, ax2]):\n",
        "        ax.set_xlabel('Turn #')\n",
        "        ax.set_ylabel('Predicted Outcome')\n",
        "        ax.grid()\n",
        "        ax.set_title(f\"Player {i+1}'s perspective\")\n",
        "    if title is not None:\n",
        "        plt.suptitle(title, y=1.2, fontsize=14)\n",
        "    plt.tight_layout()\n",
        "    plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S_v34kkghJdn",
        "cellView": "form"
      },
      "source": [
        "#@markdown Recreate ValueNet class\n",
        "\n",
        "# Set up a network for position evaluation. \n",
        "class ValueNet(nn.Module):\n",
        "    def __init__(self, num_channels=64, n=8):\n",
        "        super(ValueNet, self).__init__()\n",
        "\n",
        "        # game params\n",
        "        self.n=n\n",
        "        self.board_x, self.board_y = self.n, self.n\n",
        "        # + 1 allows for passing\n",
        "        self.action_size = self.n*self.n + 1\n",
        "        self.num_channels = num_channels\n",
        "\n",
        "        # two conv layers\n",
        "        self.conv1 = nn.Conv2d(1,  num_channels, 3, stride=1)\n",
        "        self.conv2 = nn.Conv2d(num_channels,  num_channels, 3, stride=1)\n",
        "\n",
        "        # two linear layers\n",
        "        self.fc1 = nn.Linear(num_channels*(self.board_x-4)*(self.board_y-4), num_channels*2)\n",
        "        self.fc2 = nn.Linear(num_channels*2, 1)\n",
        "\n",
        "    def forward(self, s):\n",
        "        s = s.view(-1, 1, self.board_x, self.board_y)               \n",
        "        s = F.relu(self.conv1(s))                     \n",
        "        s = F.relu(self.conv2(s))                         \n",
        "        s = s.view(-1, self.num_channels*(self.board_x-4)*(self.board_y-4))\n",
        "\n",
        "        s = F.relu(self.fc1(s))\n",
        "        v = self.fc2(s)                                                                       \n",
        "\n",
        "        return torch.tanh(v)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U856-10tP6Ik"
      },
      "source": [
        "There are two notable differences from the last notebook. First, any function that evaluates a batch of board states – a value heuristic, for example – must now also take a vector called `whoami` that indicates (for each board) which player's perspective we are evaluating from. Today not all boards will have the same player.\n",
        "\n",
        "To give an example, we've provided an example value heuristic that could have been used in last notebook. This function is like `my_board_value` from part 1 but now `whoami` is a `[n_games]` size tensor"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X7AIvlJ7P61F"
      },
      "source": [
        "def my_board_value(board_states:torch.Tensor,\n",
        "                   whoami:torch.Tensor) -> torch.Tensor:\n",
        "    \"\"\"\n",
        "    Simply returns the proportion of tiles for each player, rescaled to [-1,1]. \n",
        "    \"\"\"\n",
        "    n_games, n, _ = board_states.size()\n",
        "    values = torch.zeros(n_games)\n",
        "    \n",
        "    n1 = (board_states == OthelloGame.PLAYER1).view(n_games, -1).sum(dim=1)\n",
        "    n2 = (board_states == OthelloGame.PLAYER2).view(n_games, -1).sum(dim=1)\n",
        "\n",
        "    frac1 = n1/(n1+n2)\n",
        "    values = frac1*2-1\n",
        "    values[whoami==OthelloGame.PLAYER2] = -values[whoami==OthelloGame.PLAYER2]\n",
        "\n",
        "    return values"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5EigXk2X0nMn"
      },
      "source": [
        "The second difference is that, once we get to tree-search agents later, AI agents will store persistent state between moves. For this reason, they now include a `new_game()` method which resets the state, which is called at the start of `ai_vs_ai`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8vpGHMPkuYyf"
      },
      "source": [
        "---\n",
        "# Section 1: Policy functions\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8HCJNRxbvKP8",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Potential policies\n",
        "try: t1;\n",
        "except NameError: t1=time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"C6YWnynO9AM\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8WdTUow7vLYJ"
      },
      "source": [
        "\n",
        "The $\\epsilon$-greedy agents above all make \"snap judgments\" by looking a single move ahead and evaluating the resulting board configuration. One problem with this is that all moves are tried from the current position, even if they are \"obviously\" bad. Ideally, we could look at a board and only evaluate the moves that look initially \"sensible.\"\n",
        "\n",
        "This leads to another kind of \"snap judgment\": a policy function. A policy is a function that takes in a board and outputs a probability distribution over possible moves without first \"trying out\" any of those moves.\n",
        "\n",
        "Earlier, when you were playing Othello against classmates or against an AI, did you mentally try out every available move? Or did you only mentally simulate \"promising\" moves? What did you base that judgment on?\n",
        "\n",
        "We'll start by looking at policy functions in isolation before returing to the idea of smart *searching* using a combination of policy functions to guide the search and value functions to score them (the Monte Carlo Tree Search or MCTS algorithm)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dR1pHI9f0swo"
      },
      "source": [
        "## Exercise 1: Hand-coding a policy agent\n",
        "\n",
        "We're going to build an agent who samples moves from a policy function. The policy is a probability distribution over moves. The agent will *sample* from the policy to select a move.\n",
        "\n",
        "First, we'll build a base class that works for all policies.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uZNLuzC5vbsK"
      },
      "source": [
        "class PolicyAgentBase(ArtificialPlayer):\n",
        "    \"\"\"Base class for game-playing using a policy. Sub-classes need only provide\n",
        "    a get_policy function.\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, name=None, temperature:float=1.0):\n",
        "        self.name = self.__class__.__name__ if name is None else name\n",
        "        self.t = temperature\n",
        "    \n",
        "    def get_policy(self, game:OthelloGame) -> torch.Tensor:\n",
        "        \"\"\"PolicyAgentBase.get_policy return a distribution over the board\n",
        "        of \"good\" moves for the current player. Doesn't need to be normalized\n",
        "        over even all legal (those are handled by PolicyAgentBase.select_move).\n",
        "\n",
        "        Warning: Must not modify game state!\n",
        "        \"\"\"\n",
        "        raise NotImplementedError(\"Must be implemented by sub-class!\")\n",
        "\n",
        "    def get_selected_move_index(self, policy):\n",
        "        \"\"\"Takes a 1d flattened policy for each board and outputs the index of\n",
        "        the action selected, sampling with the appropriate temperature.\n",
        "        \n",
        "        Inputs: policy, a (n_games, n_moves) Tensor\n",
        "        Outputs: choice_idx, a (n_games) Tensor with integers specifying\n",
        "                 the action selected for each game (i.e. the column of policy)\n",
        "        \"\"\"\n",
        "        if self.t < 1e-6:\n",
        "            # Super low temperature is treated as greedy\n",
        "            choice_idx = torch.argmax(policy, dim=1, keepdim=True)\n",
        "        else:\n",
        "            choice_idx = torch.multinomial(policy**(1/self.t), num_samples=1)\n",
        "        return choice_idx    \n",
        "    \n",
        "    def select_move(self, game:OthelloGame):\n",
        "        policy = self.get_policy(game)\n",
        "        policy = policy * game._valid_moves.to(policy.device)\n",
        "        flat_policy = policy.view(game.n_games, -1)\n",
        "        # Place a dummy value into boards that have no legal moves\n",
        "        no_moves = torch.sum(game._valid_moves.view(game.n_games, -1), dim=1) == 0\n",
        "        flat_policy[no_moves, 0] = 1\n",
        "\n",
        "        flat_choice_idx = self.get_selected_move_index(flat_policy)\n",
        "\n",
        "        # Now we need to take this unrolled index and convert it into the x,y \n",
        "        # board position of the selected move\n",
        "        i_max, j_max = flat_choice_idx // game.n, flat_choice_idx % game.n\n",
        "        actions_ij = torch.cat([i_max, j_max], dim=1)\n",
        "        # Pass on all games that had no valid moves\n",
        "        actions_ij[no_moves, :] = -1\n",
        "        return actions_ij\n",
        "    \n",
        "    def plot_policy(self, game, mask_legal=True):\n",
        "        with torch.no_grad():\n",
        "            policy = self.get_policy(game).cpu()\n",
        "            if mask_legal:\n",
        "                policy = policy * game._valid_moves.cpu()\n",
        "            # Normalize it\n",
        "            policy = policy / torch.sum(policy, dim=(1,2), keepdims=True)\n",
        "\n",
        "            for g in range(game.n_games):\n",
        "                plt.figure(figsize=(6,6))\n",
        "                ax = plt.gca()\n",
        "                im = plt.imshow(policy[g,:,:])\n",
        "                plt.xticks([]); plt.yticks([])\n",
        "                suffix = '(full board)' if not mask_legal else '(legal moves only)'\n",
        "                plt.title(f'Policy heatmap {suffix}')\n",
        "                divider = make_axes_locatable(ax)\n",
        "                cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
        "                plt.colorbar(im, cax=cax)\n",
        "                plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vk7NzV8Gsd6H"
      },
      "source": [
        "Here is a very, very simple policy. It cares about all moves equally. Note that it doesn't need to return valid moves - masking out invalid moves is handled in `select_move` above. Also note that it doesn't need to normalize its probability distribution to 1. That is also handled in `select_move`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nt3_O7SqN_Q4"
      },
      "source": [
        "class UniformPolicyAgent(PolicyAgentBase):\n",
        "    def get_policy(self, game:OthelloGame) -> torch.Tensor:\n",
        "        return torch.ones_like(game.boards)\n",
        "\n",
        "game = OthelloGame(n_games=1, n=8)\n",
        "game.render()\n",
        "uniform_agent = UniformPolicyAgent()\n",
        "uniform_agent.plot_policy(game, mask_legal=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_N3L5_nTOArP"
      },
      "source": [
        "Let's imagine creating a hand-coded function for a policy for Othello, just like you created one for value last week. Talk in your group and come up with a good strategy."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IDno-btTOV5M",
        "cellView": "form"
      },
      "source": [
        "policy_strategy = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "18HphilJOozw"
      },
      "source": [
        "## Section 1.1: Making a policy agent\n",
        "\n",
        "_Time estimate: 30 minutes since start_"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U60TyyQiOo6_",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Handmade policies\n",
        "param_list = ['policy_strategy']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "video = YouTubeVideo(id=\"0pB3jdzRI0U\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ftW9W7QkOV-J"
      },
      "source": [
        "We won't make you actually code out a hand-coded policy. This is a deep learning course, after all. Instead, we've provided a heuristic for you. This agent simply likes edges and corners better than the center of the board."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0F2XNTUmsvxX",
        "collapsed": true
      },
      "source": [
        "class MyPolicyAgent(PolicyAgentBase):\n",
        "    def get_policy(self, game:OthelloGame) -> torch.Tensor:\n",
        "        \"\"\"Input: game, an instance of OthelloGame\n",
        "         Returns: policy a (n_games, n, n) Tensor. The distribution over the \n",
        "                  board of \"good\" moves for the current player. \n",
        "        \"\"\"\n",
        "\n",
        "        policy = torch.ones_like(game.boards)\n",
        "        for i in range(game.n//2):\n",
        "            policy[:, game.n//2-1-i, game.n//2-1-i:game.n//2+i] = i\n",
        "            policy[:, game.n//2+i, game.n//2-1-i:game.n//2+i+1] = i\n",
        "            policy[:, game.n//2-1-i:game.n//2+i+1, game.n//2-1-i] = i\n",
        "            policy[:, game.n//2-1-i:game.n//2+i+1, game.n//2+i] = i\n",
        "        return policy\n",
        "\n",
        "game = OthelloGame(n_games=1, n=8)\n",
        "game.step(game.get_available_actions()[0][0])\n",
        "game.step(game.get_available_actions()[0][0])\n",
        "game.render()\n",
        "my_agent = MyPolicyAgent()\n",
        "my_agent.plot_policy(game, mask_legal=False)\n",
        "my_agent.plot_policy(game, mask_legal=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YfDVFtdRnlSH"
      },
      "source": [
        "How do these policy-based snap-judgement agents compare against each other?\n",
        "\n",
        "Or against the value-based agents?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BPSriBgjnjm_"
      },
      "source": [
        "uniform_agent = UniformPolicyAgent()\n",
        "my_policy_agent = MyPolicyAgent()\n",
        "heuristic_value_agent = EpsilonGreedyAgent(name=\"GreedyValueAgent\")\n",
        "\n",
        "win_statistics(my_policy_agent, uniform_agent)\n",
        "win_statistics(my_policy_agent, heuristic_value_agent)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IkFsj2Z-ShCf"
      },
      "source": [
        "## Exercise 2: How can we train a policy network?\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KBW8Hcb1jxXG",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Training a Policy Network\n",
        "\n",
        "video = YouTubeVideo(id=\"MMzxB47l5bM\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "owWiB2hujwwF"
      },
      "source": [
        "Let's say we make a NN that takes in a board an outputs a distribution over board locations - a policy. Recall from the value network that training a NN requires \n",
        "1. a quantifiable objective to minimize, and \n",
        "2. a lot of data.\n",
        "\n",
        "Can we formulate an appropriate objective function to train towards, and build a proper labeled training set?\n",
        "\n",
        "Take 2 minutes in silence (before group discussion) and write down a strategy for training a neural network to output a good policy."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "i805GsUBLgZl",
        "cellView": "form"
      },
      "source": [
        "how_to_train_policy= '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GZ58WfbPHcym"
      },
      "source": [
        "Now read out your strategies in your pod. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BDn_Q32ZnUT9",
        "cellView": "form"
      },
      "source": [
        "#@title Video: RL and historical ways of thinking\n",
        "param_list = ['how_to_train_policy']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "video = YouTubeVideo(id=\"llgQ5EKAORg\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video\n",
        "    "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hBheB62NpSxv"
      },
      "source": [
        "---\n",
        "# Section 2: Moving beyond snap-judgments: planning and tree-search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g1a7Bwp0T_fk"
      },
      "source": [
        "_Time estimate: 50 minutes since start_"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2pXZC87wwKky",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Trees that describe the future\n",
        "\n",
        "try: t2;\n",
        "except NameError: t2 = time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"Mo2fwXqJGzM\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K1-wXtvJwDOx"
      },
      "source": [
        "Above, we saw examples of agents that play Othello entirely by \"snap judgments\". The Value-based agents looked one move ahead and evaluated that position. The Policy-based agents map directly from a board to a distribution of plausible actions without even looking a single step ahead.\n",
        "\n",
        "A much more powerful way of selecting moves is to treat the policy as a \"hunch\" about what moves are good, then simulate a few steps into the future, scoring the result a few moves down the line using the value function.\n",
        "\n",
        "Why do we need to use our policy \"hunch\" when planning, when all we are interested is the estimated value a few moves out? There are simply too many possible moves. Here we've created the move tree for just the first full turn."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mWOLwm1PqNfe",
        "cellView": "form"
      },
      "source": [
        "#@markdown MCTSTree class and visualization tools (just run me - we'll see more on MCTS later)\n",
        "import plotly.graph_objs as go\n",
        "import networkx as nx\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "\n",
        "class MCTSTree(object):\n",
        "    \"\"\"Monte Carlo Tree Search. An MCTS object is initialized with a policy and\n",
        "    value function.\n",
        "\n",
        "    Search behavior is controlled using num_search (number of searches to run)\n",
        "    and cpuct (explore/exploit hyperparameter)\n",
        "\n",
        "    Based partially on github.com/suragnair/alpha-zero-general/\n",
        "    \"\"\"\n",
        "\n",
        "    def __init__(self, policy_value_fun, n_games=1, n=8, num_search=100, cpuct:float=1.0):\n",
        "        assert cpuct > 0, \"cpuct parameter must be positive\"\n",
        "        assert num_search >= 1, \"Must have >=1 simulation\"\n",
        "        self.pol_val_fun = policy_value_fun\n",
        "        self.num_search = num_search\n",
        "        self.cpuct = cpuct\n",
        "        self.max_depth = 500\n",
        "        self.n_games = n_games\n",
        "        self.n = n\n",
        "        self.init_tree()\n",
        "\n",
        "        if self.n_games > 1 and self.__class__.__name__ == \"MCTSTree\":\n",
        "            raise ValueError(\"The base MCTSTree class can only handle one game at a time!\")\n",
        "    \n",
        "    def init_tree(self):\n",
        "        # s refers to state, a to actions. States are used as dictionary keys\n",
        "        # by calling state_key on single games.\n",
        "        self.Ns = {}  # stores #times state s was visited\n",
        "        self.Ls = {}  # stores list of legal moves at state s\n",
        "        self.Ps = {}  # stores snap policy judgment at s: tensor same size as Ls\n",
        "        self.Vs = {}  # stores snap value judgment at s: scalar tensor\n",
        "        self.Tsa = {} # stores total sub-tree values for s: tensor same size as Ls where index i corresponds to Ls[s][i]\n",
        "        self.Nsa = {} # stores #times edge s,a was visited: same format as Tsa\n",
        "        # For debugging/visualization...\n",
        "        self.Ksa = {} # stores a set of (action, child_idx) tuples from each parent node\n",
        "        self.Bs = {} # stores board state as a tuple of integers for plotting\n",
        "    \n",
        "    def state_key(self, game:OthelloGame, player):\n",
        "        # get_uid returns a [n_games] tensor of int64 hashes. Assume n_games=1\n",
        "        # and grab the int out of the tensor.\n",
        "        return game.get_uid(player).item()\n",
        "\n",
        "    def count_child_visits(self, game):\n",
        "        \"\"\"Return visit count for children of the given game\n",
        "        \"\"\"\n",
        "        s = self.state_key(game, game.current_player)\n",
        "        counts = torch.zeros_like(game.boards)\n",
        "        children = self.Ls[s]\n",
        "        for k, (i,j) in enumerate(children):\n",
        "            counts[0,i,j] = self.Nsa[s][k]\n",
        "        return counts\n",
        "    \n",
        "    def run_searches(self, game):\n",
        "        \"\"\"Run self.num_searches searches from the given game position, extending\n",
        "        whatever tree we already have.\n",
        "\n",
        "        Note: nothing is returned, since the effect of this function is to\n",
        "        expand the tree which is stored as instance variables.\n",
        "        \"\"\"\n",
        "        for i in range(self.num_search):\n",
        "            state = game.copy_state()\n",
        "            self.single_search(game)\n",
        "            game.paste_state(state)\n",
        "\n",
        "    def single_search(self, game:OthelloGame):\n",
        "        \"\"\"This function performs one iteration of MCTS. It loops until a leaf \n",
        "        state (unvisited or end of game) is found. The action chosen at each\n",
        "        step is one that has the maximum upper confidence bound (UCB).\n",
        "\n",
        "        Note: nothing is returned, since the effect of this function is to\n",
        "        expand the tree which is stored as instance variables.\n",
        "\n",
        "        After returning, assume 'game' is irreversible altered. The caller is\n",
        "        responsible for copying the game state as needed. (See @run_searches)\n",
        "        \"\"\"\n",
        "        path = []\n",
        "        for d in range(self.max_depth):\n",
        "            s = self.state_key(game, game.current_player)\n",
        "\n",
        "            # Record board position for plotting\n",
        "            if s not in self.Bs:\n",
        "                self.Bs[s] = tuple(int(v) for v in game.boards[0].flatten().cpu())\n",
        "            \n",
        "            # Count visit to this state\n",
        "            self.Ns[s] = 1 if s not in self.Ns else self.Ns[s]+1\n",
        "\n",
        "            is_unvisited = self.Ns[s] == 1\n",
        "            is_leaf = is_unvisited or game.game_status == OthelloGame.OVER\n",
        "            if is_leaf:\n",
        "                # Found a leaf! Get ready to propagate value back through the\n",
        "                # search path.\n",
        "                if is_unvisited:\n",
        "                    if game.game_status == OthelloGame.OVER:\n",
        "                        # No guesswork on the value. Use the actual game outcome.\n",
        "                        result = game.score_games()[0]\n",
        "                        self.Vs[s] = result[0] if game.current_player == OthelloGame.PLAYER1 else result[1]\n",
        "                    else:\n",
        "                        # This is a *new* leaf. Run policy and value on this state.\n",
        "                        pol, self.Vs[s] = self.pol_val_fun(game.boards, game.current_player)\n",
        "                        self.Ls[s] = game.get_available_actions()[0]\n",
        "                        # Store snap-policy judgment just over legal moves\n",
        "                        self.Ps[s] = torch.tensor([pol[0,i,j] for (i,j) in self.Ls[s]])\n",
        "                        # Initialize empty child counts from s\n",
        "                        self.Nsa[s] = torch.zeros(len(self.Ls[s]))\n",
        "                        self.Tsa[s] = torch.zeros(len(self.Ls[s]))\n",
        "                # Note: value Vs[s] is from the perspective of the player whos\n",
        "                # turn it is in state s, *not* from the perspective of whoever\n",
        "                # just played a turn before.\n",
        "                backup_value = self.Vs[s]\n",
        "                break\n",
        "            else:\n",
        "                # We've been here before. Select next move according to UCB.\n",
        "                legal_actions = self.Ls[s]\n",
        "                if len(legal_actions) == 0:\n",
        "                    next_action_idx, next_action = None, OthelloGame.PASS\n",
        "                else:\n",
        "                    # Q-value is average value = total value / number of visits\n",
        "                    Qsa = self.Tsa[s] / self.Nsa[s]\n",
        "                    # Wherever Nsa==0, set Q to (otherwise it is nan)\n",
        "                    Qsa[self.Nsa[s] == 0] = 0.0\n",
        "                    # UCB is a combination of Q values and snap-judgment policy.\n",
        "                    # Note: for UCB purposes, this visit isn't complete yet, so subtract 1 from Ns\n",
        "                    UCB = Qsa + self.cpuct * self.Ps[s] * math.sqrt(self.Ns[s] - 1) / (1 + self.Nsa[s])\n",
        "                    # Argmax the UCB score for next action\n",
        "                    next_action_idx = torch.argmax(UCB).item()\n",
        "                    next_action = self.Ls[s][next_action_idx]\n",
        "                \n",
        "                # Book-keeping: keep track of all (s,a) pairs along the search path\n",
        "                path.append((s, next_action_idx))\n",
        "\n",
        "                # Advance to the next state\n",
        "                game.step(next_action)\n",
        "                \n",
        "                # Store parent id -> (action, child id) map\n",
        "                if next_action_idx is not None:\n",
        "                    if s not in self.Ksa:\n",
        "                        self.Ksa[s] = set()\n",
        "                    self.Ksa[s].add((next_action, self.state_key(game, game.current_player)))\n",
        "        else:\n",
        "            # For...else syntax should be read as \"no break\" clause. We land\n",
        "            # here if no leaf was ever encountered.\n",
        "            raise RuntimeError(\"Never hit a leaf! This shouldn't happen!\")\n",
        "        \n",
        "        # print(\"[DEBUG] ENDED SEARCH AT\", s)\n",
        "        # print(\"[DEBUG] PATH WAS\", path)\n",
        "        \n",
        "        # Run back through the path updating states. Recall that values are always\n",
        "        # from the perspective of whoever's turn it is when the board was evaluated.\n",
        "        # Since player 1's gains are necessarily player 2's losses (AKA minimax),\n",
        "        # we have to flip the sign of the value for each step back up the tree.\n",
        "        while len(path) > 0:\n",
        "            backup_value = -backup_value\n",
        "            s, action_idx = path.pop()\n",
        "            # print(\"[DEBUG] BACKUP\", s, action_idx, self.Ls[s][action_idx])\n",
        "            # print(\"[DEBUG] VALUE\", backup_value)\n",
        "            if action_idx is not None:\n",
        "                self.Nsa[s][action_idx] = self.Nsa[s][action_idx] + 1\n",
        "                self.Tsa[s][action_idx] = self.Tsa[s][action_idx] + backup_value\n",
        "    \n",
        "    def to_graph(self, game:OthelloGame, g_idx=0, max_depth:int=100):\n",
        "        \"\"\"Debugging helper. Retuns a networkx.DiGraph representation of the tree.\n",
        "        \"\"\"\n",
        "        root = self.state_key(game, game.current_player)\n",
        "        G = nx.DiGraph()\n",
        "        G.add_node(root,\n",
        "                player=game.current_player,\n",
        "                value=self.Vs[root].item(),\n",
        "                # hash=root, # for debugging only\n",
        "                n=self.Ns[root])\n",
        "        # queue contains tuples of (parent uid, child idx, child uid, depth, player)\n",
        "        queue = [(root, act, ch, 1, 3-game.current_player) for act, ch in self.Ksa[root]]\n",
        "        while len(queue) > 0:\n",
        "            parent, act, child, depth, player = queue.pop()\n",
        "            idx = [i for i in range(len(self.Ls[parent])) if self.Ls[parent][i] == act][0]\n",
        "            G.add_node(child,\n",
        "                    player=player,\n",
        "                    value=self.Vs[child].item(),\n",
        "                    # hash=child, # for debugging only\n",
        "                    visits=self.Ns[child])\n",
        "            G.add_edge(parent, child,\n",
        "                    policy=self.Ps[parent][idx].item(),\n",
        "                    visits=self.Nsa[parent][idx].item(),\n",
        "                    q=self.Tsa[parent][idx].item()/self.Nsa[parent][idx].item(),\n",
        "                    action=act)\n",
        "            if depth < max_depth and child in self.Ksa:\n",
        "                queue.extend([(child, a, ch, depth+1, 3-player) for a, ch in self.Ksa[child]])\n",
        "        return G, root\n",
        "\n",
        "def plot_mcts_tree(tree:MCTSTree, game:OthelloGame, max_depth:int=100, edge_width_range=(0.5,3), title=None, no_mcts_metadata=False):\n",
        "    G, root = tree.to_graph(game, g_idx=0, max_depth=max_depth)\n",
        "    \n",
        "    # Based on https://plotly.com/python/network-graphs/\n",
        "\n",
        "    # Compute x,y locations for each node that lays out the plot nicely\n",
        "    layout = {}\n",
        "    parents = [root]\n",
        "    d=0\n",
        "    while len(parents) > 0:\n",
        "        children = []\n",
        "        for i, p in enumerate(parents):\n",
        "            layout[p] = (i+1)/(len(parents)+1), 1-d/max_depth\n",
        "            if p in tree.Ksa:\n",
        "                children.extend([k for _, k in tree.Ksa[p] if k is not None])\n",
        "        d += 1\n",
        "        parents = children\n",
        "    \n",
        "    min_width, max_width = edge_width_range\n",
        "\n",
        "    edge_traces, edge_hovers = [], []\n",
        "    for u,v in G.edges():\n",
        "        x0, y0 = layout[u]\n",
        "        x1, y1 = layout[v]\n",
        "        norm = max(attr['visits'] for _, attr in G[u].items())\n",
        "        edge_traces.append(go.Scatter(\n",
        "            x=[x0, x1], y=[y0, y1], mode='lines',\n",
        "            line={'width': min_width+(max_width-min_width)*G[u][v]['visits']/norm, 'color': 'black'}))\n",
        "        edge_txt = [f\"{k}={v}\" for k,v in G[u][v].items()]\n",
        "        if not no_mcts_metadata:\n",
        "          edge_hovers.append(go.Scatter(\n",
        "              x=[(x0+x1)/2], y=[(y0+y1)/2], mode='markers',\n",
        "              marker={'color':'black', 'size':5},\n",
        "              hoverinfo='text', text=\"<br>\".join(edge_txt)))\n",
        "          \n",
        "    \n",
        "    node_x, node_y, node_text, node_player = [], [], [], []\n",
        "    for node in G.nodes():\n",
        "        x, y = layout[node]\n",
        "        node_x.append(x)\n",
        "        node_y.append(y)\n",
        "        node_player.append(G.nodes[node]['player'])\n",
        "        node_txt = [f\"{k}={v}\" for k,v in G.nodes[node].items()]\n",
        "        board_string = \" \".join(OthelloGame.GAME_SYMBOLS[v] for v in tree.Bs[node]) if node in tree.Bs else \"[no board saved]\"\n",
        "        board_string_2d = \"<br>\".join(board_string[2*i:2*(i+game.n)] for i in range(0, game.n**2, game.n))\n",
        "        node_text.append(board_string_2d + \"<br>\" + \"<br>\".join(node_txt))\n",
        "        if no_mcts_metadata:\n",
        "          node_text[-1] = board_string_2d\n",
        "\n",
        "\n",
        "    node_trace = go.Scatter(x=node_x, y=node_y, text=node_text, mode='markers', hoverinfo='text',\n",
        "                            marker={'size':10,'cmin':1, 'cmax':2,\n",
        "                                    'color': node_player,\n",
        "                                    'colorscale': ['black', 'white'],\n",
        "                                    'line': {'width': 2, 'color':'gray'}})\n",
        "    \n",
        "    fig = go.Figure(data=edge_traces+edge_hovers+[node_trace],\n",
        "            layout=go.Layout(title='MCTS Tree' if title is None else title, titlefont_size=16, showlegend=False,\n",
        "                            width=500*(1 if no_mcts_metadata else 1.5), height=400*(1 if no_mcts_metadata else 1.5),\n",
        "                            hovermode='closest', hoverlabel={'bgcolor':'white'},\n",
        "                            xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),\n",
        "                            yaxis=dict(showgrid=False, zeroline=False, showticklabels=False)))\n",
        "    fig.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qRh11ENR3GD2"
      },
      "source": [
        "def naive_policy_value_fun(boards:torch.Tensor, whoami:torch.Tensor):\n",
        "    pol = torch.ones_like(boards)\n",
        "    val = torch.zeros_like(boards)[:,0,0]\n",
        "    return pol, val\n",
        "\n",
        "def plot_beginning_tree():\n",
        "    game = OthelloGame()\n",
        "    tree = MCTSTree(naive_policy_value_fun, num_search=17, cpuct=1.0)\n",
        "    tree.run_searches(game)\n",
        "    plot_mcts_tree(tree, game, title=\"Hover over markers to see boards\", no_mcts_metadata=True)\n",
        "\n",
        "plot_beginning_tree()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0mzdOyHJGBrf"
      },
      "source": [
        "**Student response:** How many final board states are there for an 8x8 board? How many valid moves exist *on average* during the course of the game?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HmUKryyZGAHC",
        "cellView": "form"
      },
      "source": [
        "final_boards = \"\" #@param {type:\"string\"}\n",
        "average_valid_moves = \"\" #@param {type:\"string\"}\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZT4LI3ClazKW"
      },
      "source": [
        "---\n",
        "# Section 3: Monte-Carlo Tree Search (MCTS): a smart way to plan\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fSvHVE0lUZY0"
      },
      "source": [
        "_Time estimate: 65 minutes since start_"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uxPAqfnW06ar",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Monte Carlo Tree Search (MCTS)\n",
        "\n",
        "param_list = ['final_boards', 'average_valid_moves']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "try: t3;\n",
        "except NameError: t3 = time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"Ibvgl1-TwQg\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UJLIartn3PcE"
      },
      "source": [
        "\n",
        "Even in Othello (let alone Go) here are far more game trees than can ever be explicitly tested. We need a way to test only promising paths. This is what MCTS does, and it is the core idea of AlphaZero. What is \"promising\" for MCTS is decided from a combination of game outcomes, value estimates, and policies. The value estimates and policies are where neural networks come in.\n",
        "\n",
        "Today's tutorial of MCTS should be a review for you. Most of you saw MCTS in CIS 521 (Machine Learning). There are also many great resources online that you should consult if today's materials are confusing (we recommend [this](https://web.archive.org/web/20180623055344/http://mcts.ai/about/index.html), [this](https://web.stanford.edu/~surag/posts/alphazero.html), and [this](http://joshvarty.github.io/AlphaZero/))."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w7AK4iyK0_d4"
      },
      "source": [
        "\n",
        "## Section 3.1: MCTS is a tree-growing algorithm\n",
        "In the beginning of each game, MCTS starts with no tree. Then, each turn, MCTS grows its tree by a set number of nodes (for us, about 100). Each node is grown by starting from the root of the tree and then following a path (i.e. simulating gameplay) with a special equation for how actions are selected. This equation is called the **upper confidence bound** (UCB). It balances the policy, value, and visit counts to balance exploration and exploitation.\n",
        "\n",
        "### The UCB\n",
        "\n",
        "Each edge of the tree $(s,a)$ maintains the following attributes:\n",
        " - $Q(s,a)$, the average value after following action $a$.\n",
        " - $N(s,a)$, the number of times this edge has been taken.\n",
        " - $N(s) = \\sum_aN(s,a)$, the number of times a node has been visited.\n",
        "\n",
        "We also store the policy $P(s,a)$, which is expensive to compute and never changes. While growing nodes, MCTS chooses actions $a$ from each node $s$ by maximizing the UCB equation:\n",
        "\n",
        "$$U(s,a)=Q(s,a) + c_{puct} P(s,a)\\frac{\\sqrt{N(s)}}{1+N(s,a)}$$\n",
        "\n",
        "$c_{puct}$ is a hyperparameter that controls exploration.\n",
        "\n",
        "When a node is expanded into new leaves, we encounter the reward (if the game ends) or an expected reward, a.k.a. a value estimate. This is progagated down to update all $Q(s,a)$ in the edges we just traversed.\n",
        "\n",
        "### A metaphor\n",
        "\n",
        "MCTS is like growing a real biological tree. Trees want to grow leaves to maximize sunlight (value). Growing leaves requires sugar in the form of sap, which needs to be pumped up from the root. How should a growing tree decide where to send sap? When a node $s$ receives sap, it sends it upwards depending on:\n",
        " 1. The policy $P(s,a)$ (an initial hunch about where sunlight is, provided in our case by a Policy Network)\n",
        " 2. The value of the next node $Q(s,a)$ (the average amount of sunlight we've seen by sending sap down this branch before, provided in our case by a Value Network)\n",
        " 3. The number of times the next node has been visited already $N(s,a)$ (branches that have received little attention in the past get some \"exploratory\" sap)\n",
        "\n",
        "When the sap reaches the final node, it grows new leaves. These leaves communicates the sunlight it received (the value) down to all of the nodes visited in this episode, which update their average value attribute ($Q$) with this new leaf value."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_PIIe6cy0_kT"
      },
      "source": [
        "## Section 3.2: Exploring tree growth"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5O5BrGLJ3blG"
      },
      "source": [
        "Before we dig into the code, let's examine correct behavior starting from the root node. We'll need a policy and value function of the boards, of course. Let's choose a simple one: a uniform policy, and a value that is the simple `my_board_value` we defined above (fraction of the total pieces that are `whoami`)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "TlDJrpjs7iAw"
      },
      "source": [
        "def random_policy_value_fun(boards:torch.Tensor, whoami:torch.Tensor):\n",
        "    pol = torch.ones_like(boards)\n",
        "    val = my_board_value(boards, whoami)\n",
        "    return pol, val"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZKiGLQOd7vpg"
      },
      "source": [
        "### The first 5 searches\n",
        "\n",
        "We've built an `MCTSTree` class for you so we can see these trees. `MCTSTree` takes a value/policy function and a number of searches – how many times we send sap up to expand a node into new leaves."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NkJrguBE3sBl"
      },
      "source": [
        "game = OthelloGame()\n",
        "tree = MCTSTree(random_policy_value_fun, num_search=5)\n",
        "tree.run_searches(game)\n",
        "plot_mcts_tree(tree, game, title=\"Hover over the edge and nodes\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sOZnWazH8XLA"
      },
      "source": [
        "Take a look at the visit count on each of the edges. Why do they sum to 4, even though we have completed 5 searches?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8LB8n1613sDn",
        "cellView": "form"
      },
      "source": [
        "sum_to_4 = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z-cw5oaCAeli"
      },
      "source": [
        "Look at the q values along the edges, and compare it to the values in the nodes (which just records the random `0.0` or `1.0` result of `random_policy_value_fun` on that node's board). Why is the sign inverted?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OJlfruV33sGA",
        "cellView": "form"
      },
      "source": [
        "sign_of_q = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GMED0bu8BjiA"
      },
      "source": [
        "### The next 5 searches"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fpuLwIdkBpcG",
        "cellView": "form"
      },
      "source": [
        "#@title Video: How Q is calculated\n",
        "\n",
        "param_list = [ 'sum_to_4', 'sign_of_q']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "video = YouTubeVideo(id=\"ByX8cr9jKu4\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E3Eyrm0o7x_3"
      },
      "source": [
        "Calling `tree.run_searches(game)` a second time continues adding to the previous tree. The number of new nodes is controlled by `tree.num_search`. The plot shows thicker lines along edges that have been visited more than their siblings in the tree.\n",
        "\n",
        "Every time you run the next cell, it will add another 5 nodes to the tree."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BGB_LBChBvBT"
      },
      "source": [
        "tree.run_searches(game)\n",
        "plot_mcts_tree(tree, game, title=\"Hover over the edge and nodes\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y8c4zKy7DJ-f"
      },
      "source": [
        "Look at the visit counts on the edges in the first row. Does it better reflect the policy, or the value estimate?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N-WfXXsnDRDa",
        "cellView": "form"
      },
      "source": [
        "visit_counts_meaning = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yH9_3Lkodrik"
      },
      "source": [
        "## Section 3.3: The effect of a policy\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "foSxbu-ugQI7"
      },
      "source": [
        "_Time estimate: 80 minutes from start_"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OIDAZtDufHwv",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Nontrivial Policies\n",
        "param_list = [ 'visit_counts_meaning']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "video = YouTubeVideo(id=\"7blA2q1pMoM\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qVr_orGPR7SU"
      },
      "source": [
        "Before, our `random_policy_value_fun` had a uniform policy, and all exploration was driven by value. Let's try the opposite. We'll use an uninformative value function (always zeros) and policy that really likes some directions over others. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sB1W9RGjSCkq"
      },
      "source": [
        "def nonuniform_policy_value_fun(boards:torch.Tensor, whoami:torch.Tensor):\n",
        "    pol = torch.bernoulli(torch.ones_like(boards)/2.) + 1e-6\n",
        "    val = torch.zeros_like(boards[:,0,0])\n",
        "    return pol, val"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FjFTIhNv44Dr"
      },
      "source": [
        "Now check out your policy in action:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FFnZYhcV5F9o"
      },
      "source": [
        "game = OthelloGame()\n",
        "tree = MCTSTree(nonuniform_policy_value_fun, num_search=30)\n",
        "tree.run_searches(game)\n",
        "plot_mcts_tree(tree, game, title=\"Hover over the edge and nodes\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1jJ7PC5x6vYE"
      },
      "source": [
        "## Section 3.4: The effect of $c_{puct}$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UZg9AIl26wUS",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Exploration Parameter\n",
        "video = YouTubeVideo(id=\"EaKMI9Su4Wc\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kh3GNmgv66Ko"
      },
      "source": [
        "Recall that while growing nodes, MCTS chooses actions $a$ from each node $s$ by maximizing the UCB equation (pasted from above):\n",
        "\n",
        "$$U(s,a)=Q(s,a) + c_{puct} P(s,a)\\frac{\\sqrt{N(s)}}{1+N(s,a)}$$\n",
        "\n",
        "Vary the value of $c_{puct}$ below. Try values close to 0 and quite large. What do you observe?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1rsqU3vV6_4o"
      },
      "source": [
        "game = OthelloGame()\n",
        "tree = MCTSTree(random_policy_value_fun, num_search=50, cpuct=10)\n",
        "tree.run_searches(game)\n",
        "plot_mcts_tree(tree, game, title=\"Hover over the edge and nodes\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-iN3FbUf7R_A",
        "cellView": "form"
      },
      "source": [
        "effect_of_cpuct = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XaYERNDBD3nz"
      },
      "source": [
        "## Section 3.5: Selecting a move for actual gameplay"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DRE1Ceh1U5i3"
      },
      "source": [
        "_Time estimate: 90 minutes since start. This next Exercise 3 should take 30 minutes for the tutorial to end on time._"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gNpazw9cEODv",
        "cellView": "form"
      },
      "source": [
        "#@title Video: How to Choose an Action\n",
        "param_list = [ 'effect_of_cpuct']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "video = YouTubeVideo(id=\"BC3cNXV-oKM\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gIQlXIBKEOLH"
      },
      "source": [
        "After running many searches, MCTS will have built up a nice tree. How should we use that tree to select a move during actual gameplay?\n",
        "\n",
        "Counterintuitively, we will **not** use the UCB equation. Otherwise, what would be the point of building up that whole exploratory tree? \n",
        "\n",
        "Instead, our agent will somehow *use* the information we gained during explanation to select the right move\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D-lonmYXi9Af"
      },
      "source": [
        "### Exercise 3.a\n",
        "\n",
        "Discuss and describe a possible strategy for using information obtained in an MCTS search tree for action selection. Keep it simple!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "n9f8zxiqi9Sd",
        "cellView": "form"
      },
      "source": [
        "action_selection_strategy = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kNeCHt_9-4Uy"
      },
      "source": [
        "### Exercise 3.b\n",
        "\n",
        "Build a subclass of `PolicyAgentBase` that uses an MCTS tree to select a move. Your agent's `get_policy` function (which is what you'll be implementing) should return a (possibly un-normalized) distribution over \"good\" moves. What \"good\" means in this context will be up to you. \n",
        "\n",
        "We're providing the MCTS tree implementation for you. An MCTS tree contains the following attributes, which are populated by calls to `run_searches`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "z60DkMG-_-T5"
      },
      "source": [
        "def init_tree(mcts):\n",
        "    # s refers to state, a to actions. States are used as dictionary keys\n",
        "    # by calling state_key on single games.\n",
        "    mcts.Ns = {}  # stores #times state s was visited\n",
        "    mcts.Ls = {}  # stores list of legal moves at state s as (i,j) index pairs\n",
        "    mcts.Ps = {}  # stores snap policy judgment at s: tensor same size as Ls where index i corresponds to Ls[s][i]\n",
        "    mcts.Vs = {}  # stores snap value judgment at s: scalar tensor\n",
        "    mcts.Tsa = {} # stores total sub-tree values for s: tensor same size as Ls where index i corresponds to Ls[s][i]\n",
        "    mcts.Nsa = {} # stores #times edge s,a was visited: same format as Tsa\n",
        "\n",
        "def state_key(game:OthelloGame, player):\n",
        "    # get_uid returns a [n_games] tensor of int64 hashes. Assume n_games=1\n",
        "    # and grab the int out of the tensor.\n",
        "    return game.get_uid(player).item()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yGU8pk8QAXPy"
      },
      "source": [
        "The dictionaries all have keys that are boards... or rather, *hashes* of those boards. The hash is provided by the `state_key` function. \n",
        "\n",
        "    # number of times the search tried out each move.\n",
        "    tree.Nsa[state_key(game, game.current_player)]  \n",
        "\n",
        "If you instead want `Q` (the _average_ value of each action taken from each state), then you might consider something like\n",
        "\n",
        "    key = state_key(game, game.current_player)\n",
        "    Q = tree.Tsa[key] / tree.Nsa[key]\n",
        "\n",
        "\n",
        "Now it's your turn. For now we are simplifying to a single game at a time since `MCTSTree` was not built to handle multiple games at once. So, the output of `get_policy` will be size `[1, n, n]`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qPzVCcg7Ag79"
      },
      "source": [
        "class MCTSAgent(PolicyAgentBase):\n",
        "    def __init__(self, tree:MCTSTree, temperature=1.0):\n",
        "        super(MCTSAgent, self).__init__(temperature=temperature)\n",
        "        self.tree = tree # you'll need this!\n",
        "    \n",
        "    def new_game(self):\n",
        "        \"\"\"Reset the tree for new games.\n",
        "        \"\"\"\n",
        "        self.tree.init_tree()\n",
        "\n",
        "    def get_policy(self, game:OthelloGame) -> torch.Tensor:\n",
        "        \"\"\"Input: game, an instance of OthelloGame\n",
        "                  MCTStree, a built tree with the root as the starting board\n",
        "         Returns: policy a (n_games, n, n) Tensor. The distribution over the \n",
        "                  board of \"good\" moves for the current player. \n",
        "        \"\"\"\n",
        "        # Run searches\n",
        "        self.tree.run_searches(game)\n",
        "        \n",
        "        #####################################################################\n",
        "        # Fill in missing code (...),\n",
        "        # then remove or comment the line below to test your function\n",
        "        #  \n",
        "        #\n",
        "        # Warning: Must not modify game state (e.g. game.boards)!\n",
        "        raise NotImplementedError(\"Complete the get_policy function of MyMCTSPolicyAgent\")\n",
        "        #####################################################################\n",
        "        \n",
        "        ...\n",
        "\n",
        "        return policy\n",
        "\n",
        "# # uncomment to test code\n",
        "# game = OthelloGame()\n",
        "# tree = MCTSTree(random_policy_value_fun, num_search=50)\n",
        "# my_mcts_agent = MCTSAgent(tree)\n",
        "# act = my_mcts_agent.select_move(game)\n",
        "# game.step(act)\n",
        "# game.render()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sTOZHeJvEARg"
      },
      "source": [
        "# to_remove solution\n",
        "class MCTSAgent(PolicyAgentBase):\n",
        "    def __init__(self, tree:MCTSTree, temperature=1.0):\n",
        "        super(MCTSAgent, self).__init__(temperature=temperature)\n",
        "        self.tree = tree # you'll need this!\n",
        "    \n",
        "    def new_game(self):\n",
        "        \"\"\"Reset the tree for new games.\n",
        "        \"\"\"\n",
        "        self.tree.init_tree()\n",
        "\n",
        "    def get_policy(self, game:OthelloGame) -> torch.Tensor:\n",
        "        \"\"\"Input: game, an instance of OthelloGame\n",
        "                  MCTStree, a built tree with the root as the starting board\n",
        "         Returns: policy a (n_games, n, n) Tensor. The distribution over the \n",
        "                  board of \"good\" moves for the current player. \n",
        "        \"\"\"\n",
        "        # Run searches\n",
        "        self.tree.run_searches(game)\n",
        "\n",
        "        # The policy is proportional to visit counts of the children\n",
        "        s = self.tree.state_key(game, game.current_player)\n",
        "        policy = torch.zeros_like(game.boards)\n",
        "        children = self.tree.Ls[s]\n",
        "        for k, (i,j) in enumerate(children):\n",
        "            policy[0,i,j] = self.tree.Nsa[s][k]\n",
        "        return policy\n",
        "\n",
        "# uncomment to test code\n",
        "game = OthelloGame()\n",
        "tree = MCTSTree(random_policy_value_fun, num_search=50)\n",
        "my_mcts_agent = MCTSAgent(tree)\n",
        "act = my_mcts_agent.select_move(game)\n",
        "game.step(act)\n",
        "game.render()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LpEKcS5tlPls"
      },
      "source": [
        "### What AlphaZero does"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p3P3homvCQV_",
        "cellView": "form"
      },
      "source": [
        "#@title Video: AlphaZero and Scaling\n",
        "param_list = ['action_selection_strategy']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "try: t4; \n",
        "except NameError: t4 = time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"fiU-d9QBQdU\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rr46ZqOEBpqg"
      },
      "source": [
        "The standard MCTS algorithm uses visit counts as the actual policy. Move selection happens simply by examining the visit counts of the next available actions. You can think of MCTS as using a simple snap-judgment policy (and a value estimate and search) to obtain a new, better policy: the visit counts. \n",
        "\n",
        "The `MCTSTree` class in fact provides a `count_child_visits` function. Here is a minimalistic `MCTSAgent` who uses `count_child_visits` only. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3ATGoBYvCIe8"
      },
      "source": [
        "class MCTSAgent(PolicyAgentBase):\n",
        "    def __init__(self, tree:MCTSTree, temperature=1.0):\n",
        "        super(MCTSAgent, self).__init__(temperature=temperature)\n",
        "        self.tree = tree # you'll need this!\n",
        "    \n",
        "    def new_game(self):\n",
        "        \"\"\"Reset the tree for new games.\n",
        "        \"\"\"\n",
        "        self.tree.init_tree()\n",
        "\n",
        "    def get_policy(self, game:OthelloGame) -> torch.Tensor:\n",
        "        \"\"\"Input: game, an instance of OthelloGame\n",
        "                  MCTStree, a built tree with the root as the starting board\n",
        "         Returns: policy a (n_games, n, n) Tensor. The distribution over the \n",
        "                  board of \"good\" moves for the current player. \n",
        "        \"\"\"\n",
        "        # Run searches\n",
        "        self.tree.run_searches(game)\n",
        "        # The policy is proportional to visit counts of the children\n",
        "        return self.tree.count_child_visits(game)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J0Teq26OWT6C"
      },
      "source": [
        "---\n",
        "# Section 4: Using MCTS to training a policy (and a value)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bYSowRotVgeQ"
      },
      "source": [
        "_Time estimate: 120 minutes since start._"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WeywcYYMWuZE"
      },
      "source": [
        "\n",
        "\n",
        "AlphaZero improves its policy network (used for tree exploration) by making it more similar to the MCTS visit counts. Let the output of a policy network with parameters $\\theta$ be $\\vec p_\\theta(s)$. Then, let the (normalized) visit counts be the vector $\\vec\\pi(s)$ with entries $\\vec \\pi(s)_i = \\frac{N(s,a_i)}{\\sum_a N(s,a)}$. The loss function we'll minimize is:\n",
        "$$ l(\\theta) = \\textrm{CrossEntropy}(\\vec\\pi(s), \\vec p_\\theta(s))\\\\ = -\\vec\\pi(s) \\cdot \\log \\vec p_\\theta(s)$$\n",
        "\n",
        "The loss function for a value network is the same as before: the mean-squared error from a value estimate to the actual outcome of the game under *real* (i.e. under $\\vec\\pi(s)$) play."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rQuV1a5JZxcY"
      },
      "source": [
        "## Exercise 4: How long will it take to train?\n",
        "\n",
        "How long does it take to build an MCTS tree? This will depend on your choice of `policy_value_fun`, undoubtely. But we can get an approximate time estimate.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "P1PMRQpjahDA"
      },
      "source": [
        "tree = MCTSTree(random_policy_value_fun, num_search=50)\n",
        "\n",
        "%timeit tree.run_searches(game)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O5-_zwtlceJ6"
      },
      "source": [
        "Try changing the number of searches. What sort of scaling do you observe?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qdAJDRi-arC4",
        "cellView": "form"
      },
      "source": [
        "search_scaling = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pBnfPPrycrBH"
      },
      "source": [
        "Now, imagine that we will be training a network for 1,000 iterations, each with a batch size of 128 games and 50 searches. (Our `MCTSTree` above only builds 1 tree at time - no parallelization). How long will that take?\n",
        "\n",
        "Again, this is a lower bound (since we are using our heuristics instead of our neural networks)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "71lYYlcjbFcm",
        "cellView": "form"
      },
      "source": [
        "train_time_bound = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q2t10wBghIl7"
      },
      "source": [
        "How long does a NN take to evaluate, anyways? Remember we defined `ValueNet()` above. How long does it take to evaluate - and how does that affect your time estimate?\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lzVLab_OyhwF"
      },
      "source": [
        "n_games = 1\n",
        "game = OthelloGame(n_games=n_games)\n",
        "value_net = ValueNet().cuda()\n",
        "\n",
        "# warm things up (pytorch does some compilation/optimization the first time\n",
        "# a function is called)\n",
        "value_net(game.boards)\n",
        "whoami = torch.ones(n_games).cuda()\n",
        "\n",
        "%timeit my_board_value(game.boards, whoami)\n",
        "%timeit value_net(game.boards)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JqsPB9SpSXuu",
        "cellView": "form"
      },
      "source": [
        "adjusted_train_time_bound = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qteOcdtDSeby"
      },
      "source": [
        "Now, as a teaser, run this plot to observe how these evaluation times change when run in parallel across boards."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "plmLwj-mjuv1",
        "cellView": "form"
      },
      "source": [
        "#@markdown Time vs. batch size (run to plot)\n",
        "timings = np.zeros((256,2))\n",
        "for n_games in range(1,256):\n",
        "    game = OthelloGame(n_games=n_games)\n",
        "    whoami = torch.ones(n_games).cuda()\n",
        "    t0 = time.time()\n",
        "    my_board_value(game.boards, whoami)\n",
        "    t1 = time.time()\n",
        "    value_net(game.boards)\n",
        "    t2 = time.time()\n",
        "    timings[n_games, :] = (t1-t0, t2-t1)\n",
        "plt.plot(timings*1000)\n",
        "plt.ylabel(\"ms to run batch of boards\")\n",
        "plt.xlabel(\"# of boards in parallel\")\n",
        "plt.legend([\"my_board_value\",\"value_net\"])\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Efb2LygtnW5S"
      },
      "source": [
        "This should convince you that computing in parallel is probably the right way to go.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bcbw5KZuyhwG"
      },
      "source": [
        "---\n",
        "# Section 5: Parallel MCTS\n",
        "\n",
        "_Time estimate: 130 minutes from start_"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "OqpGqOBpyhwH",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Running MCTS in Parallel\n",
        "param_list = ['search_scaling', 'train_time_bound', 'adjusted_train_time_bound']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "       raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "try: t5; \n",
        "except NameError: t5 = time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"P-bhG9BXPgo\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pNAKy-yBmfHw"
      },
      "source": [
        "As you calculated, it will be absolutely crucial to speed up serial tree building. What can we parallelize? What can we speed?\n",
        "\n",
        "The easiest thing to parallelize over is the games. In fact, we've been doing this all week! Below, we've built some functions that run MCTS in parallel over the boards in an `OthelloGame`.\n",
        "\n",
        "Thus far the `MCTSTree` is implemented in very readable Python. This is for your sake! However, this means that 1) it operates one board at a time, and 2) all of the attributes are located on the CPU. We'd rather run everything on a GPU. Next, we will define such an all-GPU MCTS agent.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vTNdOKvyyhwH",
        "cellView": "form"
      },
      "source": [
        "#@markdown CUDA helper kernel (just run me)\n",
        "\n",
        "@cuda.jit\n",
        "def cuda_index_of(values, table_entries, table_lengths, indices, add_new):\n",
        "    \"\"\"Naive linear search to find the index(es) of the given value(es). Runs in\n",
        "    parallel: multiple tables and values can be searched at once. Results are\n",
        "    stored in 'indices'. If add_new=True and an item isn't found, it is appended\n",
        "    to the table and the new index is returned. If add_new=False, then -1 is used\n",
        "    to mean 'not found'.\n",
        "\n",
        "    Uses block indexing, so call with cuda_index_of[num_tables,1](...)\n",
        "    \"\"\"\n",
        "    bx = cuda.blockIdx.x\n",
        "    max_search = table_lengths[bx]\n",
        "    val = values[bx]\n",
        "\n",
        "    for i in range(max_search):\n",
        "        if table_entries[bx, i] == val:\n",
        "            indices[bx] = i\n",
        "            return\n",
        "\n",
        "    # Not found! either add a new entry or return -1\n",
        "    if add_new:\n",
        "        table_entries[bx, max_search] = val\n",
        "        table_lengths[bx] = table_lengths[bx] + 1\n",
        "        indices[bx] = max_search\n",
        "    else:\n",
        "        indices[bx] = -1\n",
        "\n",
        "def test_cuda_index_of():\n",
        "    table = torch.zeros(2, 10, dtype=torch.int64, device='cuda')\n",
        "    hashes = torch.zeros(2, dtype=torch.int64, device='cuda')\n",
        "    lengths = torch.zeros(2, dtype=torch.int32, device='cuda')\n",
        "    indices = torch.zeros(2, dtype=torch.int32, device='cuda')\n",
        "\n",
        "    hashes[0] = 42\n",
        "    hashes[1] = 101\n",
        "\n",
        "    # Empty table --> results should have index -1 and table and lengths should be unchanged\n",
        "    cuda_index_of[2,1](torch2cuda(hashes), torch2cuda(table), torch2cuda(lengths), torch2cuda(indices), False)\n",
        "    assert torch.all(indices == -1)\n",
        "    assert torch.all(lengths == 0)\n",
        "    assert torch.all(table == 0)\n",
        "\n",
        "    # Same inputs except now add_new=True. Indices should be 0 (1st element) and\n",
        "    # table and lengths should be updated\n",
        "    cuda_index_of[2,1](torch2cuda(hashes), torch2cuda(table), torch2cuda(lengths), torch2cuda(indices), True)\n",
        "    assert torch.all(indices == 0)\n",
        "    assert torch.all(lengths == 1)\n",
        "    assert torch.all(table[:,0] == hashes)\n",
        "    assert torch.all(table[:,1:] == 0)\n",
        "\n",
        "    # add_new a second time with same inputs --> no effect\n",
        "    cuda_index_of[2,1](torch2cuda(hashes), torch2cuda(table), torch2cuda(lengths), torch2cuda(indices), True)\n",
        "    assert torch.all(indices == 0)\n",
        "    assert torch.all(lengths == 1)\n",
        "    assert torch.all(table[:,0] == hashes)\n",
        "    assert torch.all(table[:,1:] == 0)\n",
        "\n",
        "    # Another round of new values, add_new=False\n",
        "    hashes[0] = 11\n",
        "    hashes[1] = 2021\n",
        "    cuda_index_of[2,1](torch2cuda(hashes), torch2cuda(table), torch2cuda(lengths), torch2cuda(indices), False)\n",
        "    assert torch.all(indices == -1)\n",
        "    assert torch.all(lengths == 1)\n",
        "    assert torch.all(table[:,0] == torch.as_tensor([42, 101], device='cuda'))\n",
        "    assert torch.all(table[:,1:] == 0)\n",
        "\n",
        "    # Another round of new values, add_new=True\n",
        "    cuda_index_of[2,1](torch2cuda(hashes), torch2cuda(table), torch2cuda(lengths), torch2cuda(indices), True)\n",
        "    assert torch.all(indices == 1)\n",
        "    assert torch.all(lengths == 2)\n",
        "    assert torch.all(table[:,0] == torch.as_tensor([42, 101], device='cuda'))\n",
        "    assert torch.all(table[:,1] == torch.as_tensor([11, 2021], device='cuda'))\n",
        "    assert torch.all(table[:,2:] == 0)\n",
        "\n",
        "    # add_new again --> no effect\n",
        "    cuda_index_of[2,1](torch2cuda(hashes), torch2cuda(table), torch2cuda(lengths), torch2cuda(indices), True)\n",
        "    assert torch.all(indices == 1)\n",
        "    assert torch.all(lengths == 2)\n",
        "    assert torch.all(table[:,0] == torch.as_tensor([42, 101], device='cuda'))\n",
        "    assert torch.all(table[:,1] == torch.as_tensor([11, 2021], device='cuda'))\n",
        "    assert torch.all(table[:,2:] == 0)\n",
        "\n",
        "    # One new one old, add_new=False\n",
        "    hashes[0] = 12345\n",
        "    hashes[1] = 101\n",
        "    cuda_index_of[2,1](torch2cuda(hashes), torch2cuda(table), torch2cuda(lengths), torch2cuda(indices), False)\n",
        "    assert torch.all(indices == torch.as_tensor([-1, 0], device='cuda'))\n",
        "    assert torch.all(lengths == torch.as_tensor([2, 2], device='cuda'))\n",
        "    assert torch.all(table[:,:2] == torch.as_tensor([[42,11],[101,2021]], device='cuda'))\n",
        "    assert torch.all(table[:,2:] == 0)\n",
        "\n",
        "    # One new one old, add_new=True\n",
        "    cuda_index_of[2,1](torch2cuda(hashes), torch2cuda(table), torch2cuda(lengths), torch2cuda(indices), True)\n",
        "    assert torch.all(indices == torch.as_tensor([2, 0], device='cuda'))\n",
        "    assert torch.all(lengths == torch.as_tensor([3, 2], device='cuda'))\n",
        "    assert torch.all(table[:,:3] == torch.as_tensor([[42,11,12345],[101,2021,0]], device='cuda'))\n",
        "    assert torch.all(table[:,3:] == 0)\n",
        "\n",
        "    # add_new again --> no effect\n",
        "    cuda_index_of[2,1](torch2cuda(hashes), torch2cuda(table), torch2cuda(lengths), torch2cuda(indices), True)\n",
        "    assert torch.all(indices == torch.as_tensor([2, 0], device='cuda'))\n",
        "    assert torch.all(lengths == torch.as_tensor([3, 2], device='cuda'))\n",
        "    assert torch.all(table[:,:3] == torch.as_tensor([[42,11,12345],[101,2021,0]], device='cuda'))\n",
        "    assert torch.all(table[:,3:] == 0)\n",
        "\n",
        "    print(\"TEST PASSED\")\n",
        "\n",
        "test_cuda_index_of()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VyubxMriyhwH",
        "cellView": "form"
      },
      "source": [
        "#@markdown Define the ParallelMCTSTree class\n",
        "class ParallelMCTSTree(MCTSTree):\n",
        "    \"\"\"Same function as MCTSTree, but parallelized over games, and searches\n",
        "    live entirely on the GPU so there is no GPU<->CPU transfer bottleneck.\n",
        "\n",
        "    policy_value_fun must also be able to accept and return entirely GPU tensors.\n",
        "\n",
        "    Overrides init_tree, count_child_visits, single_search, to_graph, and\n",
        "    state_key, but inherits run_searches and is otherwise functionally identical\n",
        "    \"\"\"\n",
        "    INITIAL_TREE_SIZE = 256\n",
        "\n",
        "    def __init__(self, *args, **kwargs):\n",
        "        if 'debug' in kwargs:\n",
        "            self._debug = kwargs['debug']\n",
        "            del kwargs['debug']\n",
        "        else:\n",
        "            self._debug = False\n",
        "        super(ParallelMCTSTree, self).__init__(*args, **kwargs)\n",
        "        # Helper for slicing into the n_games dimension, since writing\n",
        "        # torch.arange(self.n_games) over and over again is ugly.\n",
        "        self._g_idx = torch.arange(self.n_games)\n",
        "    \n",
        "    def state_key(self, game:OthelloGame, player):\n",
        "        # Whereas MCTSTree returns the .item() of a single game hash, here\n",
        "        # we want all of them. Calling get_uid with out=self.cur_hash deposits\n",
        "        # the result in the memory already owned by self.cur_hash and\n",
        "        # self._cuda_cur_hash. The 'return' is redundant; it returns self.cur_hash\n",
        "        # A call to _set_idx_to_hash(...) is still required to line up cur_idx\n",
        "        # with cur_hash!\n",
        "        return game.get_uid(player, out=self.cur_hash)\n",
        "    \n",
        "    def _set_idx_to_hash(self, add_new):\n",
        "        \"\"\"Set self.cur_idx to the index in the lookup_hash table that matches\n",
        "        self.cur_hash. If add_new is True, creates a new table entry for unseen\n",
        "        hashes. If add_new is false, then cur_idx will be -1 if unseen.\n",
        "        \"\"\"\n",
        "        cuda_index_of[self.n_games, 1](self._cuda_cur_hash,\n",
        "                                       self._cuda_lookup_hash,\n",
        "                                       self._cuda_num_nodes,\n",
        "                                       self._cuda_cur_idx,\n",
        "                                       add_new)\n",
        "\n",
        "    def init_tree(self):\n",
        "        \"\"\"Override MCTSTree.init_tree. Rather than python dicts, we do all\n",
        "        tree book-keeping using GPU tensors.\n",
        "\n",
        "        The idea is to pre-allocate space for a small-ish tree. As we search, if\n",
        "        the data structures for the tree fill up, we double its allocated space\n",
        "        and copy data over (see _expand_tree and _expand_tree_if_full)\n",
        "        \"\"\"\n",
        "        # Clear the GPU memory before trying to allocate big tensors. Otherwise,\n",
        "        # GPU memory fills up and we get esoteric error messages about \"CUDA\n",
        "        # device-side assertion failure\"\n",
        "        torch.cuda.empty_cache()\n",
        "        # Allocate a tree with space for INITIAL_TREE_SIZE nodes. During search,\n",
        "        # we will expand the memory as needed.\n",
        "        self._tree_size = ParallelMCTSTree.INITIAL_TREE_SIZE\n",
        "        self.Ns  = torch.zeros(self.n_games, self._tree_size, dtype=torch.float32, device='cuda') # One count per game per node\n",
        "        self.Vs  = torch.zeros(self.n_games, self._tree_size, dtype=torch.float32, device='cuda') # One value per game per node\n",
        "        self.Ps  = torch.zeros(self.n_games, self._tree_size, self.n, self.n, dtype=torch.float32, device='cuda') # Per-game, per-node, snap judgment policy output\n",
        "        self.Tsa = torch.zeros(self.n_games, self._tree_size, self.n, self.n, dtype=torch.float32, device='cuda') # Per-game, per-node, per-position, sum of sub-tree values\n",
        "        self.Nsa = torch.zeros(self.n_games, self._tree_size, self.n, self.n, dtype=torch.int32, device='cuda') # Per-game, per-node, per-position, visit counts\n",
        "        self.lookup_hash = torch.zeros(self.n_games, self._tree_size, dtype=torch.int64, device='cuda') # Lookup table of board hashes\n",
        "        self.cur_hash = torch.zeros(self.n_games, dtype=torch.int64, device='cuda') # Helper during search: current hash\n",
        "        self.cur_idx = torch.zeros(self.n_games, dtype=torch.int64, device='cuda') # Helper during search: current index corresponding to cur_hash\n",
        "        self.num_nodes = torch.zeros(self.n_games, dtype=torch.int32, device='cuda') # How many nodes we have so far to cap the search through lookup_hash\n",
        "        if self._debug:\n",
        "            # Only while in debug mode: keep a dictionary of parent->child UIDs\n",
        "            # just like in the MCTSTree parent class. This requires GPU->CPU\n",
        "            # data transfer so it is disabled by default. But it is necessary\n",
        "            # for to_graph().\n",
        "            self.Ksa = {}\n",
        "            # Also keep a record of uid -> board state for visualization\n",
        "            self.Bs = {}\n",
        "\n",
        "        # Make _cuda_X wrappers around variables involved in hash lookups\n",
        "        to_wrap = ['lookup_hash', 'cur_hash', 'cur_idx', 'num_nodes']\n",
        "        for var in to_wrap:\n",
        "            self.__dict__['_cuda_' + var] = torch2cuda(self.__dict__[var])\n",
        "    \n",
        "    def _expand_tree_if_full(self):\n",
        "        buffer = 2 # Give some wiggle room just in case\n",
        "        if self.num_nodes.max() >= self._tree_size - buffer:\n",
        "            self._expand_tree()\n",
        "    \n",
        "    def _expand_tree(self):\n",
        "        # print(f\"[DEBUG] resizing tree from {self._tree_size} to {self._tree_size*2}\")\n",
        "        old_size = self._tree_size\n",
        "        self._tree_size *= 2\n",
        "\n",
        "        # Make a copy of each of the to_copy variables, expanding them along\n",
        "        # dimension 1 to twice their previous size.\n",
        "        to_copy = ['Ns', 'Vs', 'Ps', 'Tsa', 'Nsa', 'lookup_hash']\n",
        "        for var in to_copy:\n",
        "            old_tensor = self.__dict__[var]\n",
        "            sz = list(old_tensor.size())\n",
        "            assert sz[1] == old_size, f\"Expected {var} to have dimension 1 size {self._tree_size} but was {sz[1]}. What gives?\"\n",
        "            sz[1] = self._tree_size\n",
        "            new_tensor = torch.zeros(tuple(sz), dtype=old_tensor.dtype, device=old_tensor.device)\n",
        "            new_tensor[:, :old_size, ...] = old_tensor\n",
        "            self.__dict__[var] = new_tensor\n",
        "        \n",
        "        # Note we didn't touch cur_hash, cur_idx, or num_nodes so those can stay as they were\n",
        "        to_wrap = ['lookup_hash']\n",
        "        for var in to_wrap:\n",
        "            self.__dict__['_cuda_' + var] = torch2cuda(self.__dict__[var])\n",
        "    \n",
        "    def count_child_visits(self, game:OthelloGame) -> torch.Tensor:\n",
        "        # First, get the hash and index of the given game\n",
        "        self.state_key(game, game.current_player)\n",
        "        self._set_idx_to_hash(add_new=False)\n",
        "        # Sanity check that all states do in fact exist in the table\n",
        "        if torch.any(self.cur_idx == -1):\n",
        "            raise ValueError(\"Failed to find a node! This shouldn't happen!\")\n",
        "        # Once index_of completes, we have the index into the nodes table per\n",
        "        # game stored in self.cur_idx\n",
        "        return self.Nsa[self._g_idx, self.cur_idx, :, :]\n",
        "\n",
        "    def single_search(self, game:OthelloGame):\n",
        "        # Path logs the (state, action) pairs along search paths. State is a\n",
        "        # copy of cur_idx, and action is a [n_games,2] tensor. Both live on GPU.\n",
        "        path = []\n",
        "\n",
        "        # Each search (probably) adds a node. Expand the tree's memory if it's\n",
        "        # close to filling up\n",
        "        self._expand_tree_if_full()\n",
        "\n",
        "        # Keep selecting moves until every game hits a leaf. We flag a game as\n",
        "        # being at a leaf by setting game_status=OVER, even though it might not\n",
        "        # really be over.\n",
        "\n",
        "        # Pre-allocate actions tensor (init to zero so the did_pass test is\n",
        "        # False on the 1st loop iteration)\n",
        "        actions = torch.zeros(self.n_games, 2, dtype=torch.int32, device='cuda')\n",
        "        # Pre-allocate boolean mask indicating which games need \"expansion\" aka\n",
        "        # need the policy_value_function to be called on them. We don't expand\n",
        "        # end-of-game states after their first visit.\n",
        "        needs_expansion = torch.zeros(self.n_games, dtype=torch.bool, device='cuda')\n",
        "        # Keep track of which game states are OVER for real (because leaf states\n",
        "        # are flagged as OVER without really being over)\n",
        "        is_terminal = torch.zeros(self.n_games, dtype=torch.bool, device='cuda')\n",
        "        # Different games may hit leaves at different times, meaning they may\n",
        "        # have different 'current_player' values! leaf_player keeps track of who\n",
        "        # 'current_player' was each time we hit a leaf. Note: if player A plays\n",
        "        # a move that results in a novel state (a leaf), then leaf_player will\n",
        "        # be player B because that's whose turn it is _in the state being\n",
        "        # evaluated_. So leaf_player is set to player B. Similarly, if player A\n",
        "        # plays the final PASS move that ends the game, then player B is again \n",
        "        # the leaf_player.\n",
        "        leaf_player = torch.zeros(self.n_games, dtype=torch.int64, device='cuda')\n",
        "        depth = 0\n",
        "        while depth < self.max_depth and torch.any(leaf_player == 0):\n",
        "            depth += 1\n",
        "\n",
        "            # Compute hashes for each game's current state -> result stored in\n",
        "            # self.cur_hash. \n",
        "            self.state_key(game, game.current_player)\n",
        "            # Lookup node index by hash -> result stored in self.cur_idx. Set\n",
        "            # add_new to True to add this state to the table if it didn't\n",
        "            # already exist\n",
        "            self._set_idx_to_hash(add_new=True)\n",
        "            # print(f\"[DEBUG] Depth {depth} hash {self.cur_hash.cpu().numpy()}\")\n",
        "            # print(f\"[DEBUG] Depth {depth} index {self.cur_idx.cpu().numpy()}\")\n",
        "\n",
        "            # print(\"[DEBUG]\" + str(depth) + \"  \"*depth + f\"TREE SIZE {self.num_nodes.max().item()}/{self._tree_size}; INDEX {self.cur_idx.max().item()}\")\n",
        "            if self.cur_idx.max() >= self._tree_size-1:\n",
        "                raise RuntimeError(\"TREE IS FULL! This should not happen!\")\n",
        "\n",
        "            # Book-keeping for plotting\n",
        "            if self._debug:\n",
        "                for g in range(self.n_games):\n",
        "                    s = self.cur_hash[g].item()\n",
        "                    self.Bs[s] = tuple(game.boards[g].cpu().int().flatten().numpy())\n",
        "\n",
        "            #########################\n",
        "            ### DETECT NEW LEAVES ###\n",
        "            #########################\n",
        "\n",
        "            # Did we hit any new leaves? If so, mark them as leaves that need\n",
        "            # expanding. (This does apply to end-of-game states that we're\n",
        "            # seeing for the first time)\n",
        "            needs_expansion = needs_expansion | (self.Ns[self._g_idx, self.cur_idx] == 0)\n",
        "            # We can tell that this leaf was *just* encountered because it will\n",
        "            # being the state of both needs_expansion but also ACTIVE. Record\n",
        "            # who the player is at the time of encountering this leaf.\n",
        "            is_new_leaf = needs_expansion & (game.game_status == OthelloGame.ACTIVE)\n",
        "            # Flag all new leaves as OVER to freeze their board state.\n",
        "            game.game_status[is_new_leaf] = OthelloGame.OVER\n",
        "            # Record leaf_player: whose turn is it *now* (\"player B\" in the\n",
        "            # comments above)\n",
        "            leaf_player[is_new_leaf] = game.current_player\n",
        "\n",
        "            # Detect games that *just* ended and store whose turn it is *now*\n",
        "            # (\"player B\" in above comments)\n",
        "            just_ended = (game.game_status == OthelloGame.OVER) & (leaf_player == 0)\n",
        "            # Mark this game as in a \"terminal\" state\n",
        "            is_terminal = is_terminal | just_ended\n",
        "            # Record whose turn it is\n",
        "            leaf_player[just_ended] = game.current_player\n",
        "\n",
        "            ########################\n",
        "            ### INCREMENT VISITS ###\n",
        "            ########################\n",
        "\n",
        "            # Increment visit count unless the previous move was a PASS, in which\n",
        "            # case we're seeing the same state over and over again but it is not\n",
        "            # a new visit.\n",
        "            did_pass = actions[:,0] == -1\n",
        "            self.Ns[self._g_idx[~did_pass], self.cur_idx[~did_pass]] = \\\n",
        "                self.Ns[self._g_idx[~did_pass], self.cur_idx[~did_pass]] + 1\n",
        "\n",
        "            # We can skip the rest of the loop if everything is done. This is\n",
        "            # important to ensure that path.append() was called exactly as many\n",
        "            # times as game.step().\n",
        "            if torch.all(game.game_status == OthelloGame.OVER):\n",
        "                # print(f\"[DEBUG] BREAKING on {depth}\")\n",
        "                break\n",
        "\n",
        "            ############################\n",
        "            ### UCB ACTION SELECTION ###\n",
        "            ############################\n",
        "            \n",
        "            # For all non-leaf games, select another move using max of UCB.\n",
        "            T = self.Tsa[self._g_idx, self.cur_idx, ...]\n",
        "            N = self.Nsa[self._g_idx, self.cur_idx, ...]\n",
        "            P = self.Ps[self._g_idx, self.cur_idx, ...]\n",
        "            Ntot = self.Ns[self._g_idx, self.cur_idx].view(self.n_games,1,1) - 1 # For UCB purposes, this visit isn't complete yet, so subtract 1\n",
        "            # Q-value is average value = total value / number of visits, or zero\n",
        "            # if the (s,a) pair has not yet been tried. Shape is [n_games, n, n]\n",
        "            Q = T/N\n",
        "            Q[N==0] = 0\n",
        "            # UCB is a combination of Q values and snap-judgment policy. Shape\n",
        "            # is [n_games, n, n]\n",
        "            UCB = Q + self.cpuct * P * torch.sqrt(Ntot) / (1 + N)\n",
        "            UCB.masked_fill_(game._valid_moves == 0, float('-inf'))\n",
        "            # Argmax the UCB score for next action\n",
        "            next_action_flat_idx = torch.argmax(UCB.view(self.n_games, -1), dim=1)\n",
        "            actions[:,0], actions[:,1] = next_action_flat_idx // self.n, next_action_flat_idx % self.n\n",
        "\n",
        "            # Pass if there were no moves available or if the game is OVER. This\n",
        "            # includes games that are flagged as leaves.\n",
        "            do_pass = torch.all(game._valid_moves.view(self.n_games, -1) == 0, dim=1) | (game.game_status == OthelloGame.OVER)\n",
        "            actions[do_pass, :] = -1\n",
        "\n",
        "            ##############################\n",
        "            ### RECORD ACTION AND STEP ###\n",
        "            ##############################\n",
        "            \n",
        "            # Book-keeping: keep track of all (state idx, action) tuples along\n",
        "            # the search path so we can back-up values later.\n",
        "            path.append((self.cur_idx.clone(), actions.clone()))\n",
        "\n",
        "            if self._debug:\n",
        "                # If in debug mode, make a copy of the 'parent' hash value\n",
        "                parent_hash = self.cur_hash.cpu().numpy()\n",
        "\n",
        "            # print(f\"[DEBUG] Depth {depth} actions are {actions.cpu().numpy()}\")\n",
        "\n",
        "            # Advance to the next state.\n",
        "            game.step(actions)\n",
        "\n",
        "            # When in debug mode, keep track of set of parent->child hashes\n",
        "            if self._debug:\n",
        "                child_hash = game.get_uid().cpu().numpy()\n",
        "                actions_copy = actions.cpu().int().numpy()\n",
        "                # Record parent hash -> set of tuples of (action, child hashes)\n",
        "                for g in range(self.n_games):\n",
        "                    s = self.cur_hash[g].item()\n",
        "                    # Get action as a tuple of (row,col) integers\n",
        "                    act = tuple(actions_copy[g,:].flatten())\n",
        "                    if do_pass[g]:\n",
        "                        continue\n",
        "                    elif parent_hash[g] in self.Ksa:\n",
        "                        self.Ksa[s].add((act, child_hash[g]))\n",
        "                    else:\n",
        "                        self.Ksa[s] = {(act, child_hash[g])}\n",
        "\n",
        "        # After loop: ensure hashes and indices are up to date with game states.\n",
        "        # Importantly, state_key depends on who the player is. Make sure all\n",
        "        # hashes are relative to the leaf_player!\n",
        "        self.state_key(game, leaf_player)\n",
        "        self._set_idx_to_hash(add_new=True)\n",
        "        # print(f\"[DEBUG] After loop, hashes are {self.cur_hash}\")\n",
        "        # print(f\"[DEBUG] ...and indices are {self.cur_idx}\")\n",
        "\n",
        "        # Sanity check that there is a leaf_player for all games\n",
        "        # print(f\"[DEBUG] leaf_player = {leaf_player.cpu().numpy()}\")\n",
        "        assert torch.all(leaf_player > 0), \"Missing leaf player! This shouldn't happen!\"\n",
        "\n",
        "        # Prepare back-up values, which will all be from the perspective of\n",
        "        # 'leaf_player'. Begin with the final-outcome value which will be used\n",
        "        # for terminal states.\n",
        "        backup_value = game.score_games()[self._g_idx, leaf_player-1].float()\n",
        "        \n",
        "        # \"Expand\" states that are being visited for the first time. This means\n",
        "        # calling the policy+value function.\n",
        "        needs_expansion = needs_expansion & ~is_terminal # don't bother expanding terminal states\n",
        "        if torch.any(needs_expansion):\n",
        "            pol_expansion, val_expansion = self.pol_val_fun(game.boards[needs_expansion, ...], leaf_player[needs_expansion])\n",
        "\n",
        "            # Store new expansion results\n",
        "            self.Vs[self._g_idx[needs_expansion], self.cur_idx[needs_expansion]] = val_expansion\n",
        "            self.Ps[self._g_idx[needs_expansion], self.cur_idx[needs_expansion], ...] = pol_expansion\n",
        "            \n",
        "            # Use estimated value for expanded states (note that terminal\n",
        "            # states' values are left unchanged)\n",
        "            backup_value[needs_expansion] = val_expansion\n",
        "\n",
        "        # Which player did we end on? Note leaf_player may be different for\n",
        "        # each game, but backup assumes all values are w.r.t. the current player.\n",
        "        # We address this by flipping all value signs wherever leaf_player is\n",
        "        # not equal to the current_player. The backup routine does not\n",
        "        # record values for PASS moves, and all games that hit a leaf early on\n",
        "        # were padded with PASS actions. By aligning all the values to the\n",
        "        # perspective of the \"current_player\", the backup is synchronized to\n",
        "        # the same player perspective across all games. This works becaues all\n",
        "        # games will go through the same number of backup steps; those that hit\n",
        "        # a leaf early will simply back-up through multiple PASSes, which has\n",
        "        # no effect.\n",
        "        backup_value[game.current_player != leaf_player] = -backup_value[game.current_player != leaf_player]\n",
        "\n",
        "        # Run back through the path updating states. Recall that values are always\n",
        "        # from the perspective of whoever's turn it is when the board was evaluated.\n",
        "        # Since player 1's gains are necessarily player 2's losses (AKA minimax),\n",
        "        # we have to flip the sign of the value for each step back up the tree.\n",
        "        while len(path) > 0:\n",
        "            node_idx, action = path.pop()\n",
        "            act_i, act_j = action[:,0].long(), action[:,1].long()\n",
        "            valid = act_i != -1 # Don't back up anything through pass actions\n",
        "            key = (self._g_idx[valid], node_idx[valid], act_i[valid], act_j[valid])\n",
        "            \n",
        "            # Calling pop() moved us one step back in time. Flip perspective.\n",
        "            backup_value = -backup_value\n",
        "\n",
        "            # Increment visit count at (state, action) pair if not pass\n",
        "            self.Nsa[key] = self.Nsa[key] + 1\n",
        "\n",
        "            # Count total value if not pass\n",
        "            self.Tsa[key] = self.Tsa[key] + backup_value[valid]\n",
        "    \n",
        "    def to_graph(self, game:OthelloGame, g_idx=0, max_depth:int=100):\n",
        "        \"\"\"Debugging helper. Retuns a networkx.DiGraph representation of the\n",
        "        tree containing useful info in the node and edge attributes.\n",
        "        \"\"\"\n",
        "        if not self._debug:\n",
        "            raise RuntimeError(\"Cannot vall to_graph if _debug is False!\")\n",
        "        root = self.state_key(game, game.current_player)[g_idx].item()\n",
        "        self._set_idx_to_hash(add_new=False)\n",
        "        G = nx.DiGraph()\n",
        "        G.add_node(root,\n",
        "                player=game.current_player,\n",
        "                value=self.Vs[g_idx, self.cur_idx[g_idx]].item(),\n",
        "                visits=self.Ns[g_idx, self.cur_idx[g_idx]])\n",
        "        # queue contains tuples of (parent uid, child idx, child uid, depth, player)\n",
        "        queue = [(root, act, ch, 1, 3-game.current_player) for act, ch in self.Ksa[root]]\n",
        "        while len(queue) > 0:\n",
        "            parent, act, child, depth, player = queue.pop()\n",
        "            self.cur_hash[:] = child\n",
        "            self._set_idx_to_hash(add_new=False)\n",
        "            if self.cur_idx[g_idx] == -1:\n",
        "                # Child not found (must be terminal)\n",
        "                continue\n",
        "            G.add_node(child,\n",
        "                    player=player,\n",
        "                    value=self.Vs[g_idx, self.cur_idx[g_idx]].item(),\n",
        "                    visits=self.Ns[g_idx, self.cur_idx[g_idx]].item())\n",
        "            child_idx = torch.where(self.lookup_hash[g_idx, ...] == child)[0].cpu().item()\n",
        "            G.add_edge(parent, child,\n",
        "                       action=act,\n",
        "                       policy=self.Ps[g_idx, child_idx, act[0], act[1]].item(),\n",
        "                       visits=self.Ns[g_idx, child_idx].item())\n",
        "            if depth < max_depth and child in self.Ksa:\n",
        "                queue.extend([(child, act, ch, depth+1, 3-player) for act, ch in self.Ksa[child]])\n",
        "        return G, root"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rZwT1ABLyhwH"
      },
      "source": [
        "# Make a graph again\n",
        "game = OthelloGame(n_games=128, n=8)\n",
        "tree = ParallelMCTSTree(random_policy_value_fun, n_games=128, n=8,\n",
        "                        num_search=100, cpuct=1.0, debug=True)\n",
        "tree.run_searches(game)\n",
        "plot_mcts_tree(tree, game) # Make plot for the 0th game"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g1xe90HRyhwI"
      },
      "source": [
        "Efficiency test: how does the Parallel MCTS agent scale with the number of games being played at once?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Yh8x4DSoyhwI",
        "cellView": "form"
      },
      "source": [
        "#@markdown Plot timings with batch size and search depth (run to plot)\n",
        "n_games = np.arange(5)*20+1 # Need >= 1 game\n",
        "n_search = np.arange(10)*10+2 # Need >= 2 searches\n",
        "gg, ss = np.meshgrid(n_games, n_search)\n",
        "par_times = np.zeros(gg.shape)\n",
        "py_times = np.zeros(n_search.shape)\n",
        "pbar = tqdm(desc='#search', total=n_search.sum())\n",
        "for i, ns in enumerate(n_search):\n",
        "    # Parallel games\n",
        "    for j, ng in enumerate(n_games):\n",
        "        game = OthelloGame(n_games=ng, n=8)\n",
        "        tree = ParallelMCTSTree(policy_value_fun=random_policy_value_fun, n_games=ng, n=8, num_search=ns)\n",
        "        tstart = time.time()\n",
        "        tree.run_searches(game)\n",
        "        par_times[i,j] = time.time() - tstart\n",
        "    \n",
        "    # Compare to pure python version\n",
        "    game = OthelloGame(n_games=1, n=8)\n",
        "    tree = MCTSTree(policy_value_fun=random_policy_value_fun, n_games=1, n=8, num_search=ns)\n",
        "    tstart = time.time()\n",
        "    tree.run_searches(game)\n",
        "    py_times[i] = time.time() - tstart\n",
        "\n",
        "    pbar.update(ns)\n",
        "pbar.close()\n",
        "\n",
        "plt.figure(figsize=(6,4))\n",
        "plt.plot(n_search, py_times, 'k')\n",
        "plt.plot(n_search, 2*py_times, '--k')\n",
        "plt.plot(n_search, 3*py_times, ':k')\n",
        "plt.plot(ss, par_times)\n",
        "plt.legend([\"Serial 1-game\", \"Est. Ser. 2-game\", \"Est. Ser. 3-game\"] + [f'{ng} // games' for ng in n_games])\n",
        "plt.xlabel('Number of searches')\n",
        "plt.ylabel('Time to 1st move (s)')\n",
        "plt.grid()\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vn8F-1Ccvusm"
      },
      "source": [
        "Describe why it is easy to parallelize over games but hard to parallelize over the searches within the MCTS tree rollout (i.e. the simulations in which actions are selected by the UCB)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3lONwHChsJ8S",
        "cellView": "form"
      },
      "source": [
        "why_tree_search_parallelization_hard = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uqDupQqve-9m"
      },
      "source": [
        "---\n",
        "# Section 6: Using MCTS to train a policy and value network\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZcZidY_6d_ff",
        "cellView": "form"
      },
      "source": [
        "#@title Video: AlphaZero - Putting it All Together\n",
        "param_list = ['why_tree_search_parallelization_hard']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "try: t6; \n",
        "except NameError: t6 = time.time()\n",
        "\n",
        "video = YouTubeVideo(id=\"xpYaV7UBsnI\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ENMWbRBPsXCA"
      },
      "source": [
        "Now that we have a working parallel MCTS environment, we can finally train a policy. Remember, the \"trick\" of AlphaZero is to use the _visit counts_ returned by MCTS as a training target for the NN, while at the same time using the NN to guide the tree-search.\n",
        "\n",
        "First, let's define a network that outputs a policy. In fact, let's define a two-headed network that will output both a policy and a value estimate. This is what AlphaZero does, too."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gO990dHXeyIs"
      },
      "source": [
        "class PolicyValueNet(nn.Module):\n",
        "    \"\"\"This is a single neural network with two outputs: one \"policy\" output and\n",
        "    one \"value\" output. It has two \"heads\" that share the same \"body\".\n",
        "    \"\"\"\n",
        "    def __init__(self, num_channels=64, n=8, dropout=0.3):\n",
        "        super(PolicyValueNet, self).__init__()\n",
        "        \n",
        "        # game params\n",
        "        self.n = n\n",
        "        self.dropout = dropout\n",
        "        self.action_size = self.n*self.n\n",
        "        self.num_channels = num_channels\n",
        "\n",
        "        # Shared body: two conv layers followed by 2 fully connected layers\n",
        "        self.conv1 = nn.Conv2d(1, num_channels, 3)\n",
        "        self.conv2 = nn.Conv2d(num_channels, num_channels*2, 3)\n",
        "        self.fc1 = nn.Linear(num_channels*2*(self.n-4)*(self.n-4), num_channels*2)\n",
        "        self.fc2 = nn.Linear(num_channels*2, num_channels)\n",
        "\n",
        "        # Value head: one more linear layer after fc2\n",
        "        self.val_fc = nn.Linear(num_channels, 1)\n",
        "\n",
        "        # Policy head: linear from fc2 to action_size\n",
        "        self.pol_fc = nn.Linear(num_channels, self.action_size)\n",
        "\n",
        "    def forward(self, s):\n",
        "        # Body\n",
        "        s = s.view(-1, 1, self.n, self.n)\n",
        "        s = F.dropout2d(F.relu(self.conv1(s)), p=self.dropout, training=self.training)\n",
        "        s = F.dropout2d(F.relu(self.conv2(s)), p=self.dropout, training=self.training)\n",
        "        s = s.view(-1, self.num_channels*2*(self.n-4)*(self.n-4))\n",
        "        s = F.dropout(F.relu(self.fc1(s)), p=self.dropout, training=self.training)\n",
        "        s = F.dropout(F.relu(self.fc2(s)), p=self.dropout, training=self.training)\n",
        "        \n",
        "        # Value head\n",
        "        v = torch.tanh(self.val_fc(s))\n",
        "        \n",
        "        # Policy head\n",
        "        p = F.log_softmax(self.pol_fc(s), dim=1)\n",
        "\n",
        "        return p, v"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ANzAprMshv0w"
      },
      "source": [
        "## Section 6.1: Preparing training examples"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IpDeR5CXDbXq"
      },
      "source": [
        "Now, our strategy is to observe self-play games of an MCTS agent and train on the following data:\n",
        " - For the policy, we need to pair the boards (the input) with the MCTS visit counts when that board was observed (the target)\n",
        " - For the value, we need to pair all boards seen in a game (input) with the eventual outcome (target)\n",
        "\n",
        "Our strategy is to give `ai_vs_ai` a callback that makes it return the visit counts on every move (inside `ai_vs_ai`, the callback is done after the `MCTSAgent` has decided on their move, but before `game.step`)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qqbn2Z4z1Tax"
      },
      "source": [
        "def mcts_visits_callback(game:OthelloGame, agents:List[MCTSAgent]):\n",
        "    \"\"\"A callback to store visit counts in the MCTS agent's tree.\n",
        "\n",
        "    Assuming both agents[0] and agents[1] are instances of MCTSAgent.\n",
        "    \"\"\"\n",
        "    player_agent = agents[game.current_player-1]\n",
        "    return player_agent.tree.count_child_visits(game).clone()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fngwTGAz1TgB"
      },
      "source": [
        "When we call `ai_vs_ai` with this callback it will return the key variables we need: `outcomes`, `raw_states`, and `raw_mcts_counts`. Just like last time, `outcomes` is a size `[n_games,2]` tensor that contains end-of-game values for player 1 in the 0th column and for player 2 in the 1st column.\n",
        "```\n",
        "outcomes, raw_states, raw_mcts_counts = ai_vs_ai(agent, callback=mcts_visits_callback)\n",
        "```\n",
        "And, just like last time, we will further process these raw states. That's what this next helper function `prepare_alpha_zero_examples` does. Don't worry too much about the implementation, but know that it:\n",
        "1. Concatenates a list of board Tensors into one big Tensor\n",
        "2. Gets the outcomes (1 per game) into the proper shape to act as targets for the value net (1 per board) \n",
        "3. Flips the sign of all boards and values that were from player 2's perspective to \"pretend\" like they are examples from player 1's perspective\n",
        "4. Normalizes visit counts so they act as probabilities\n",
        "5. Deals with passed moves\n",
        "6. Adds board symmetries (rotations, reflections, etc.) for data augmentation\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KntZyEXz01MG",
        "cellView": "form"
      },
      "source": [
        "#@markdown RUN ME: Define `prepare_alpha_zero_examples` and `MCTSagent_from_net` helpers\n",
        "\n",
        "def prepare_alpha_zero_examples(raw_states:List[torch.Tensor], \n",
        "                                raw_counts:List[torch.Tensor], \n",
        "                                player1_values:torch.Tensor):\n",
        "    \"\"\"Convert from ai_vs_ai output to training examples.\n",
        "\n",
        "    This function performs canonical preprocessing (e.g. +1 is \"current player\"\n",
        "    and -1 is \"other player\"), concatenates things, and adds in all symmetries.\n",
        "    \"\"\"\n",
        "    n_moves = len(raw_states)\n",
        "    n_games, n, _ = raw_states[0].size()\n",
        "    # Concatenate and preprocess board states. Result is size [n_games*n_moves,n,n]\n",
        "    # ordered like [game0turn0, game1turn0, ... gameNturn0, game0turn1, game1turn1, ... ]\n",
        "    nn_states = boards2nn(raw_states)\n",
        "    # Repeat values for every turn (player1_values begins as size [n_games])\n",
        "    values = player1_values.flatten().float().repeat(n_moves)\n",
        "    # Flip the sign of all boards that were player 2\n",
        "    turn_number = torch.arange(n_moves)\n",
        "    player = (turn_number % 2) + 1\n",
        "    # 'repeat_interleave()' takes inputs [a b c] and repeats them like\n",
        "    # [a a a ... b b b ... c c c ...]. Not to be confused with 'repeat()' which\n",
        "    # outputs [a b c a b c a b c ...].\n",
        "    player = player.repeat_interleave(n_games)\n",
        "    nn_states[player == 2, ...] = -nn_states[player == 2, ...]\n",
        "    values[player == 2, ...] = -values[player == 2, ...]\n",
        "    # Concatenate and normalize all MCTS visit probabilities\n",
        "    counts = torch.cat(raw_counts, dim=0).reshape(n_moves*n_games, n, n)\n",
        "    total_visits = counts.sum(dim=2, keepdim=True).sum(dim=1, keepdim=True)\n",
        "    probs = counts / total_visits\n",
        "    # It may be the case that not all games ran for n_moves. These are indicated\n",
        "    # by nan values in the board state (see ai_vs_ai). Also drop states in which\n",
        "    # there were no legal moves, indicated by probs being nan (from 0/0)\n",
        "    drop_states = torch.isnan(nn_states[:,0,0]) | torch.isnan(probs[:,0,0])\n",
        "    train_states = nn_states[~drop_states, ...]\n",
        "    train_probs  = probs[~drop_states, ...]\n",
        "    train_values = values[~drop_states]\n",
        "    # Add all symmetries (increases effective data 8-fold)\n",
        "    train_states = torch.cat([sym(train_states) for sym in SYMMETRIES], dim=0)\n",
        "    train_probs  = torch.cat([sym(train_probs) for sym in SYMMETRIES], dim=0)\n",
        "    train_values = train_values.repeat(len(SYMMETRIES))\n",
        "    \n",
        "    return train_states, train_probs, train_values\n",
        "\n",
        "def MCTSagent_from_net(policy_value_net: PolicyValueNet, n=8, \n",
        "                       games_per_iter=128, num_mcts_search=50):\n",
        "    \"\"\"A helper function that takes a PolicyValueNet and returns an MCTSAgent\n",
        "    that uses that network.\"\"\"\n",
        "\n",
        "    def pol_val_fn(boards:torch.Tensor, whoami:torch.Tensor):\n",
        "        \"\"\" This helper function creates prepares the boards to give to the network\n",
        "        and puts the output in the correct shapes. For use inside the MCTS agent.\"\"\"\n",
        "        policy_value_net.eval()\n",
        "        # Standard board to nn preprocessing: player 2 is now -1\n",
        "        boards_nn = boards2nn(boards)\n",
        "        # Always evaluate with +1 meaning \"myself\" and -1 meaning \"other player\"\n",
        "        boards_nn[whoami == OthelloGame.PLAYER2, ...] = -boards_nn[whoami == OthelloGame.PLAYER2, ...]\n",
        "        with torch.no_grad():   # saves on overhead\n",
        "            log_policy_output, value = policy_value_net(boards_nn)\n",
        "        policy = torch.exp(log_policy_output).view(-1, n, n)\n",
        "        return policy, value.flatten()\n",
        "\n",
        "    tree = ParallelMCTSTree(pol_val_fn, n_games=games_per_iter, n=n,\n",
        "                            num_search=num_mcts_search)\n",
        "    agent = MCTSAgent(tree, temperature=1.0)\n",
        "    return agent"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tEnjKhp2uUKW"
      },
      "source": [
        "Let's think about point #3. Why not just feed in raw states to the neural nets? What would go wrong?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "W4fRkZm5VuSk",
        "cellView": "form"
      },
      "source": [
        "why_flip_perspective = '' #@param {type:\"string\"}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "44koidJ0hz7F"
      },
      "source": [
        "## Section 6.2: The AlphaZero training loop"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rGXzxqZ0biEY",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Training Loop\n",
        "param_list = ['why_flip_perspective']\n",
        "for param in param_list:\n",
        "    if param not in locals():\n",
        "        raise NameError(\"Please make sure to run the cell \"\n",
        "                        \"after submitting your answer for <{}>!\".format(param))\n",
        "\n",
        "video = YouTubeVideo(id=\"OArpVKM3B9I\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vuligAL04GDM"
      },
      "source": [
        "Finally, let's build the loss function for policy. Remember this is the cross entropy:\n",
        "$$CE(\\vec\\pi(s), p_\\theta(s)) = -\\vec\\pi(s) \\cdot \\log \\vec p_\\theta(s)$$\n",
        "\n",
        "The cross entropy is minimized when $\\vec\\pi(s)=p_\\theta(s)$. Normally we might be able to use `torch.nn.CrossEntropyLoss`, but for technical reasons we have to roll our own here:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4IRPcc314Xvy"
      },
      "source": [
        "def cross_entropy_loss(target_p, log_q):\n",
        "    \"\"\"Computs the cross entropy. Assumes target_p is a normalized distribution\n",
        "    and log_q is a log normalized distribution (e.g. output of log_softmax)\"\"\"\n",
        "    b = target_p.size()[0]\n",
        "    plogq = target_p.view(b, -1) * log_q.view(b, -1)\n",
        "    cross_entropy = -plogq.sum(dim=1)\n",
        "    return cross_entropy.mean()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_0YbB7ub5G-m"
      },
      "source": [
        "Now, finally, we can define the training loop for AlphaZero!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Gu2WxVQy4XyO"
      },
      "source": [
        "def train_alpha_zero(policy_value_net, n_iters=20, steps_per_iter=50, n=8,\n",
        "                     games_per_iter=128, num_mcts_search=50):\n",
        "    \"\"\"This master funtion creates the MCTSAgent from a policy/value net,\n",
        "    then trains the networks via self-play\"\"\"\n",
        "    \n",
        "    # first create the agent\n",
        "    agent = MCTSagent_from_net(policy_value_net, n, games_per_iter, \n",
        "                               num_mcts_search)\n",
        "    \n",
        "    # create the optimizer and MSE loss function for value part\n",
        "    mse_loss = nn.MSELoss(reduction='mean')\n",
        "    opt = torch.optim.Adam(policy_value_net.parameters(), lr=1e-3)\n",
        "\n",
        "    # Training loop\n",
        "    losses = torch.zeros(n_iters*steps_per_iter).cuda()\n",
        "    pol_losses = torch.zeros(n_iters*steps_per_iter).cuda()\n",
        "    val_losses = torch.zeros(n_iters*steps_per_iter).cuda()\n",
        "    for i in tqdm(range(n_iters), desc='epochs'):\n",
        "        ##### Here we build our training examples #####\n",
        "\n",
        "        # Play AI against itself, running games_per_iter games all in parallel\n",
        "        outcomes, raw_states, raw_mcts_counts = \\\n",
        "            ai_vs_ai(agent, n=n, n_games=games_per_iter,\n",
        "                     callback=mcts_visits_callback, del_progbar=True)\n",
        "\n",
        "        # Get all trainable information from this batch of games\n",
        "        train_states, train_probs, train_values = \\\n",
        "            prepare_alpha_zero_examples(raw_states, raw_mcts_counts, outcomes[:,0])\n",
        "\n",
        "        ##### Here we train the network #####\n",
        "        for j in range(steps_per_iter):\n",
        "            policy_value_net.train() # Ensure we are in 'training' mode rather than 'evaluation' mode\n",
        "            opt.zero_grad() # ready the optimizer\n",
        "            log_policy_output, value_output = policy_value_net(train_states)\n",
        "            pol_loss = cross_entropy_loss(train_probs, log_policy_output) # policy part of loss\n",
        "            val_loss = mse_loss(value_output.flatten(), train_values) # value part of loss\n",
        "            loss = pol_loss + val_loss # total loss\n",
        "            loss.backward() # get gradients of loss w/r/t the network parameters\n",
        "            opt.step() # take a step to lower the loss\n",
        "\n",
        "            losses[i*steps_per_iter + j] = loss.detach()\n",
        "            pol_losses[i*steps_per_iter + j] = pol_loss.detach()\n",
        "            val_losses[i*steps_per_iter + j] = val_loss.detach()\n",
        "\n",
        "    return losses.cpu().numpy(), pol_losses.cpu().numpy(), val_losses.cpu().numpy()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kGyu65u3iCUa"
      },
      "source": [
        "## Section 6.3: Training the AlphaZero model"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pXJVTbkBiCaF",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Training AlphaZero\n",
        "\n",
        "video = YouTubeVideo(id=\"9p-QUnUpNEI\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eToT_nu1nI4N"
      },
      "source": [
        "Run a few iterations of the following training code (about 1 minute) and take note of the total estimated training time. If you have time, you can let it run to completion to get your very own AlphaZero Othello model.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4PvEXBT6RGuO"
      },
      "source": [
        "# Run this cell to train a model from scratch\n",
        "n=8\n",
        "pol_val_net = PolicyValueNet(n=n).cuda()\n",
        "l, p_l, v_l = train_alpha_zero(pol_val_net, n_iters=100, steps_per_iter=50, n=n,\n",
        "                               games_per_iter=128, num_mcts_search=20)\n",
        "\n",
        "## Optional - save the trained model to your local machine. Note: you can use\n",
        "## files.upload() in the future to upload your local files to colab\n",
        "# from google.colab import files\n",
        "# data = {\"state_dict\": pol_val_net.state_dict(),\n",
        "#         \"losses\":l, \"pol_losses\": p_l, \"val_losses\": v_l}\n",
        "# torch.save(data, f\"my_policy_value_weights_{n}x{n}.pt\")\n",
        "# files.download(f\"my_policy_value_weights_{n}x{n}.pt\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gBZDjEvYpJnx",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Loading Pretrained Results\n",
        "\n",
        "video = YouTubeVideo(id=\"KzWGb0ay8mM\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S67YVSW8pCj2"
      },
      "source": [
        "Here we have \"pre-trained\" a model for you which you will download in the cell following this one."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WiU0eezpoT1E"
      },
      "source": [
        "# Run this cell to download a pre-trained model. These models were trained with\n",
        "# steps_per_iter=50, games_per_iter=128, and num_mcts_search=20. For 6x6 boards\n",
        "# we trained for n_iters=100 and for 8x8 boards n_iters=200.\n",
        "if n == 6:\n",
        "    !if [ ! -f pretrained_policy_value_weights_6x6.pt ]; then wget https://osf.io/5hy6q/download -O pretrained_policy_value_weights_6x6.pt; fi\n",
        "    data = torch.load(\"pretrained_policy_value_weights_6x6.pt\")\n",
        "elif n == 8:\n",
        "    !if [ ! -f  pretrained_policy_value_weights_8x8.pt ]; then wget https://osf.io/8c6dx/download -O pretrained_policy_value_weights_8x8.pt; fi\n",
        "    data = torch.load(\"pretrained_policy_value_weights_8x8.pt\")\n",
        "\n",
        "# Load pre-trained weights into the model\n",
        "pol_val_net = PolicyValueNet(n=n).cuda()\n",
        "pol_val_net.load_state_dict(data[\"state_dict\"])\n",
        "# Load losses for plotting\n",
        "l, p_l, v_l = data[\"losses\"], data[\"pol_losses\"], data[\"val_losses\"]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YfqhBt_1o3MO"
      },
      "source": [
        "# Plot training loss from whichever model was used above\n",
        "plt.plot(l)\n",
        "plt.plot(p_l)\n",
        "plt.plot(v_l)\n",
        "plt.legend([\"Total Loss\", \"Policy Cross-Entropy\", \"Value MSE\"])\n",
        "plt.xlabel('Number of gradient steps')\n",
        "plt.xticks(np.arange(0,len(l),len(l)//10))\n",
        "plt.ylabel('Losses')\n",
        "plt.title(f'Alpha Zero Training Loss for {n}x{n} Othello')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BF1fYSL7nRgu"
      },
      "source": [
        "\n",
        "How might you evaluate if it worked? One possible idea is to see how good is the learned policy at predicting MCTS visit counts:\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fC2gXPxx1I9_",
        "cellView": "form"
      },
      "source": [
        "#@markdown (RUN ME) See alignment of policy to MCTS visit counts in a test game\n",
        "pol_val_net.eval()\n",
        "alpha_zero_agent = MCTSagent_from_net(pol_val_net, n, 100, 20)\n",
        "alpha_zero_agent.temperature = 0.1\n",
        "outcomes, raw_states, raw_mcts_counts = ai_vs_ai(alpha_zero_agent, n_games=100,\n",
        "                                                 n=n, callback=mcts_visits_callback)\n",
        "plt.figure(figsize=(10,5))\n",
        "ax1, ax2 = plt.subplot(1,2,1), plt.subplot(1,2,2)\n",
        "for t in range(len(raw_states)):\n",
        "    sgn = +1 if t % 2 == 0 else -1\n",
        "    player = 1 if t % 2 == 0 else 2\n",
        "    pol_pred, val_pred = pol_val_net(sgn*boards2nn(raw_states[t]))\n",
        "    counts_t = raw_mcts_counts[t].view(-1, n*n)\n",
        "    targets_t = counts_t / counts_t.sum(dim=1, keepdim=True)\n",
        "    valid = ~torch.any(torch.isnan(raw_states[t][:,0,0]))\n",
        "    ax1.scatter(targets_t[valid,...].cpu().T, \n",
        "                pol_pred[valid,...].detach().exp().view(-1, n*n).cpu().T, \n",
        "                marker='.', c = 'k', alpha=.1)\n",
        "    for g in range(len(raw_states)):\n",
        "        if not torch.isnan(raw_states[t][g,0,0]):\n",
        "            ax2.plot(t, val_pred[g].item(), marker='.', \n",
        "                      color=outcome2color(outcomes[g,player-1], 0.25))\n",
        "  \n",
        "ax1.set_xlabel('normalized MCTS visit counts')\n",
        "ax1.set_ylabel('policy net output')\n",
        "ax1.set_title('Policy Net Predictions')\n",
        "ax2.set_xlabel('Turn Number')\n",
        "ax2.set_ylabel('Value Net Output')\n",
        "ax2.set_title('Value Net Predictions')\n",
        "plt.show()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-08cIGGzlxoi"
      },
      "source": [
        "### (If time) play a game against your trained agent"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mODlAguUlyGQ"
      },
      "source": [
        "alpha_zero_agent = MCTSagent_from_net(pol_val_net, n=8, games_per_iter=1,\n",
        "                                      num_mcts_search=50)\n",
        "alpha_zero_agent.temperature = 0.1\n",
        "interface = InteractiveOthelloGame(player1='human',\n",
        "                                   player2=alpha_zero_agent, n=n)\n",
        "interface.next_turn()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7UFrR7a7o2aH"
      },
      "source": [
        "---\n",
        "# Wrap-up: AlphaZero and society\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "btU-9QABpFZ5",
        "cellView": "form"
      },
      "source": [
        "#@title Video: Society and Wrap Up\n",
        "\n",
        "video = YouTubeVideo(id=\"tP_80oRvufc\", width=854, height=480, fs=1)\n",
        "print(\"Video available at https://youtube.com/watch?v=\" + video.id)\n",
        "\n",
        "video"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OuU83dImqPpp"
      },
      "source": [
        "## Submit responses"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6VhCpyuUqWTX",
        "cellView": "form"
      },
      "source": [
        "#@markdown #Run Cell to Show Airtable Form\n",
        "#@markdown ##**Confirm your answers and then click \"Submit\"**\n",
        "\n",
        "import time\n",
        "import numpy as np\n",
        "from IPython.display import IFrame\n",
        "\n",
        "def prefill_form(src, fields: dict):\n",
        "  '''\n",
        "  src: the original src url to embed the form\n",
        "  fields: a dictionary of field:value pairs,\n",
        "  e.g. {\"pennkey\": my_pennkey, \"location\": my_location}\n",
        "  '''\n",
        "  prefills = \"&\".join([\"prefill_%s=%s\"%(key, fields[key]) for key in fields])\n",
        "  src = src + prefills\n",
        "  src = \"+\".join(src.split(\" \"))\n",
        "  return src\n",
        "\n",
        "\n",
        "#autofill time if it is not present\n",
        "try: t0;\n",
        "except NameError: t0 = time.time()\n",
        "try: t1;\n",
        "except NameError: t1 = time.time()\n",
        "try: t2;\n",
        "except NameError: t2 = time.time()\n",
        "try: t3;\n",
        "except NameError: t3 = time.time()\n",
        "try: t4;\n",
        "except NameError: t4 = time.time()\n",
        "try: t5;\n",
        "except NameError: t5 = time.time()\n",
        "try: t6;\n",
        "except NameError: t6 = time.time()\n",
        "try: t7;\n",
        "except NameError: t7 = time.time()\n",
        "\n",
        "#autofill fields if they are not present\n",
        "#a missing pennkey and pod will result in an Airtable warning\n",
        "#which is easily fixed user-side.\n",
        "try: my_pennkey;\n",
        "except NameError: my_pennkey = \"\"\n",
        "try: my_pod;\n",
        "except NameError: my_pod = \"Select\"\n",
        "try: policy_strategy;\n",
        "except NameError: policy_strategy = \"\"\n",
        "try: how_to_train_policy;\n",
        "except NameError: how_to_train_policy = \"\"\n",
        "try: final_boards;\n",
        "except NameError: final_boards = 0\n",
        "try: average_valid_moves;\n",
        "except NameError: average_valid_moves = 0\n",
        "try: sum_to_4;\n",
        "except NameError: sum_to_4 = \"\"\n",
        "try: sign_of_q;\n",
        "except NameError: sign_of_q = \"\"\n",
        "try: visit_counts_meaning;\n",
        "except NameError: visit_counts_meaning = \"\"\n",
        "try: effect_of_cpuct;\n",
        "except NameError: effect_of_cpuct = \"\"\n",
        "try: action_selection_strategy;\n",
        "except NameError: action_selection_strategy = \"\"\n",
        "try: search_scaling;\n",
        "except NameError: search_scaling = \"\"\n",
        "try: train_time_bound;\n",
        "except NameError: train_time_bound = \"\"\n",
        "try: adjusted_train_time_bound;\n",
        "except NameError: adjusted_train_time_bound = \"\"\n",
        "try: why_tree_search_parallelization_hard;\n",
        "except NameError: why_tree_search_parallelization_hard = \"\"\n",
        "try: why_flip_perspective;\n",
        "except NameError: why_flip_perspective = \"\"\n",
        "\n",
        "times = [(t-t0) for t in [t1,t2,t3,t4,t5,t6,t7]]\n",
        "\n",
        "fields = {\"pennkey\": my_pennkey,\n",
        "          \"pod\": my_pod,\n",
        "          \"policy_strategy\": policy_strategy,\n",
        "          \"how_to_train_policy\": how_to_train_policy,\n",
        "          \"final_boards\": final_boards,\n",
        "          \"average_valid_moves\": average_valid_moves,\n",
        "          \"sum_to_4\": sum_to_4,\n",
        "          \"sign_of_q\": sign_of_q,\n",
        "          \"visit_counts_meaning\": visit_counts_meaning,\n",
        "          \"effect_of_cpuct\": effect_of_cpuct,\n",
        "          \"action_selection_strategy\": action_selection_strategy,\n",
        "          \"search_scaling\": search_scaling,\n",
        "          \"train_time_bound\": train_time_bound,\n",
        "          \"adjusted_train_time_bound\": adjusted_train_time_bound,\n",
        "          \"why_tree_search_parallelization_hard\": why_tree_search_parallelization_hard,\n",
        "          \"why_flip_perspective\": why_flip_perspective,\n",
        "          \"cumulative_times\": times}\n",
        "src = \"https://airtable.com/embed/shrxxokvcPAFaPvj0?\"\n",
        "\n",
        "#now instead of the original source url, we do: src = prefill_form(src, fields)\n",
        "display(IFrame(src = prefill_form(src, fields), width = 800, height = 400))\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QNFRudPafVrD"
      },
      "source": [
        "## Feedback\n",
        "How could this session have been better? How happy are you in your group? How do you feel right now?\n",
        "\n",
        "Feel free to use the embeded form below or use this link:\n",
        "<a target=\"_blank\" rel=\"noopener noreferrer\" href=\"https://airtable.com/shrNSJ5ECXhNhsYss\">https://airtable.com/shrNSJ5ECXhNhsYss</a>"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "t6F7dpf5fVwC"
      },
      "source": [
        "display(IFrame(src=\"https://airtable.com/embed/shrNSJ5ECXhNhsYss?backgroundColor=red\", width = 800, height = 400))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v8FSL7cBfVEz"
      },
      "source": [
        "## Homework: \n",
        "\n",
        "Complete the following assignment: https://colab.research.google.com/github/CIS-522/course-content/blob/main/tutorials/W1_AlphaZero/student/W1_Homework.ipynb\n",
        "\n",
        "This is due 1 week from today at the start of next week's Tutorial 2."
      ]
    }
  ]
}